Merge pull request #482 from watany-dev/add-model-cl35-lama31

Add support for Claude3.5-Sonnet
aws-cloudformation · Aug 5, 2024 · e7962cf · e7962cf
2 parents cf74051 + 4acd826
commit e7962cf
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 28 deletions.
diff --git a/docs/rain_build.md b/docs/rain_build.md
@@ -18,7 +18,7 @@ rain build [<resource type>] or <prompt>
   -h, --help                 help for build
   -j, --json                 Output the template as JSON (default format: YAML)
   -l, --list                 List all CloudFormation resource types with an optional name prefix
-      --model string         The ID of the Bedrock model to use for --prompt. Shorthand: claude2, claude3haiku, claude3sonnet, claude3opus (default "claude2")
+      --model string         The ID of the Bedrock model to use for --prompt. Shorthand: claude2, claude3haiku, claude3sonnet, claude3opus, claude3.5sonnet, (default "claude2")
       --no-cache             Do not used cached schema files
       --omit-patches         Omit patches and use the raw schema
   -o, --output string        Output to a file

diff --git a/internal/cmd/build/build.go b/internal/cmd/build/build.go
@@ -436,6 +436,7 @@ func init() {
 	models["claude3opus"] = "anthropic.claude-3-opus-20240229-v1:0"
 	models["claude3sonnet"] = "anthropic.claude-3-sonnet-20240229-v1:0"
 	models["claude3haiku"] = "anthropic.claude-3-haiku-20240307-v1:0"
+	models["claude3.5sonnet"] = "anthropic.claude-3-5-sonnet-20240620-v1:0"
 
 	activeFormat = " {{ .Name | magenta }}: {{ .Text | magenta }}"
 	selectedFormat = " {{ .Name | magenta }}: {{ .Text | blue }}"
@@ -457,5 +458,5 @@ func init() {
 	Cmd.Flags().BoolVar(&pklClass, "pkl-class", false, "Output a pkl class based on a resource type schema")
 	Cmd.Flags().BoolVar(&noCache, "no-cache", false, "Do not used cached schema files")
 	Cmd.Flags().StringVar(&promptLanguage, "prompt-lang", "cfn", "The language to target for --prompt, CloudFormation YAML (cfn), CloudFormation Guard (guard), Open Policy Agent Rego (rego)")
-	Cmd.Flags().StringVar(&model, "model", "claude2", "The ID of the Bedrock model to use for --prompt. Shorthand: claude2, claude3haiku, claude3sonnet, claude3opus")
+	Cmd.Flags().StringVar(&model, "model", "claude2", "The ID of the Bedrock model to use for --prompt. Shorthand: claude2, claude3haiku, claude3sonnet, claude3opus, claude3.5sonnet")
 }
diff --git a/internal/cmd/build/interactive.go b/internal/cmd/build/interactive.go
@@ -103,9 +103,10 @@ func template() {
 	case PROMPT:
 		selections = []buildSelection{
 			{Name: CLAUDE2, Text: "Claude 2"},
-			{Name: HAIKU, Text: "Claude 3 Haiku"},
-			{Name: SONNET, Text: "Claude 3 Sonnet"},
-			{Name: OPUS, Text: "Claude 3 Opus"},
+			{Name: CLAUDE3HAIKU, Text: "Claude 3 Haiku"},
+			{Name: CLAUDE3SONNET, Text: "Claude 3 Sonnet"},
+			{Name: CLAUDE3OPUS, Text: "Claude 3 Opus"},
+			{Name: CLAUDE3_5SONNET, Text: "Claude 3.5 Sonnet"},
 		}
 		selected = openPrompt("Select a model", selections)
 		prompt := promptui.Prompt{
@@ -130,9 +131,10 @@ func policy() {
 	}
 	lang := openPrompt("Choose a language", selections)
 	selections = []buildSelection{
-		{Name: HAIKU, Text: "Claude 3 Haiku"},
-		{Name: SONNET, Text: "Claude 3 Sonnet"},
-		{Name: OPUS, Text: "Claude 3 Opus"},
+		{Name: CLAUDE3HAIKU, Text: "Claude 3 Haiku"},
+		{Name: CLAUDE3SONNET, Text: "Claude 3 Sonnet"},
+		{Name: CLAUDE3OPUS, Text: "Claude 3 Opus"},
+		{Name: CLAUDE3_5SONNET, Text: "Claude 3.5 Sonnet"},
 	}
 	model = openPrompt("Choose a model", selections)
 	prompt := promptui.Prompt{
@@ -153,9 +155,10 @@ func policy() {
 func generic() {
 	// Choose the model
 	selections := []buildSelection{
-		{Name: HAIKU, Text: "Claude 3 Haiku"},
-		{Name: SONNET, Text: "Claude 3 Sonnet"},
-		{Name: OPUS, Text: "Claude 3 Opus"},
+		{Name: CLAUDE3HAIKU, Text: "Claude 3 Haiku"},
+		{Name: CLAUDE3SONNET, Text: "Claude 3 Sonnet"},
+		{Name: CLAUDE3OPUS, Text: "Claude 3 Opus"},
+		{Name: CLAUDE3_5SONNET, Text: "Claude 3.5 Sonnet"},
 	}
 	model = openPrompt("Choose a model", selections)
 
@@ -240,6 +243,7 @@ func interactive() {
 				ii. Claude 3 Haiku
 				iii. Claude 3 Sonnet
 				iv. Claude 3 Opus
+				iv. Claude 3.5 Sonnet
 			d. Output a vetted recommended template for a use case
 		3. Create a policy validation file
 		    a. Guard
@@ -285,21 +289,22 @@ type buildSelection struct {
 }
 
 const (
-	LIST      = "list"
-	SCHEMA    = "schema"
-	TEMPLATE  = "template"
-	POLICY    = "policy"
-	ALL       = "all"
-	REQUIRED  = "required"
-	RECOMMEND = "recommend"
-	PROMPT    = "prompt"
-	CLAUDE2   = "claude2"
-	OPUS      = "claude3opus"
-	SONNET    = "claude3sonnet"
-	HAIKU     = "claude3haiku"
-	GUARD     = "guard"
-	OPA       = "opa"
-	GENERIC   = "generic"
-	YES       = "yes"
-	NO        = "no"
+	LIST               = "list"
+	SCHEMA             = "schema"
+	TEMPLATE           = "template"
+	POLICY             = "policy"
+	ALL                = "all"
+	REQUIRED           = "required"
+	RECOMMEND          = "recommend"
+	PROMPT             = "prompt"
+	CLAUDE2            = "claude2"
+	CLAUDE3OPUS        = "claude3opus"
+	CLAUDE3SONNET      = "claude3sonnet"
+	CLAUDE3HAIKU       = "claude3haiku"
+	CLAUDE3_5SONNET    = "claude3.5sonnet"
+	GUARD              = "guard"
+	OPA                = "opa"
+	GENERIC            = "generic"
+	YES                = "yes"
+	NO                 = "no"
 )