From b48fe334859c013149eaf540ada7a89ff50753eb Mon Sep 17 00:00:00 2001 From: Su Wei Date: Thu, 19 Oct 2023 17:03:58 +0800 Subject: [PATCH] add amazonbedrock, change model list to const var Signed-off-by: Su Wei --- pkg/ai/amzonbedrock.go | 17 ++++++++++++++--- 1 file changed, 14 insertions(+), 3 deletions(-) diff --git a/pkg/ai/amzonbedrock.go b/pkg/ai/amzonbedrock.go index 473efa4..a2f61c1 100644 --- a/pkg/ai/amzonbedrock.go +++ b/pkg/ai/amzonbedrock.go @@ -33,19 +33,30 @@ type InvokeModelResponseBody struct { const BEDROCK_REGION = "us-east-1" // default use us-east-1 region +const ( + ModelAnthropicClaudeV2 = "anthropic.claude-v2" + ModelAnthropicClaudeV1 = "anthropic.claude-v1" + ModelAnthropicClaudeInstantV1 = "anthropic.claude-instant-v1" +) + +var BEDROCK_MODELS = []string{ + ModelAnthropicClaudeV2, + ModelAnthropicClaudeV1, + ModelAnthropicClaudeInstantV1, +} + // GetModelOrDefault check config model func GetModelOrDefault(model string) string { - modelList := []string{"anthropic.claude-v2", "anthropic.claude-v1", "anthropic.claude-instant-v1"} // Check if the provided model is in the list - for _, m := range modelList { + for _, m := range BEDROCK_MODELS { if m == model { return model // Return the provided model } } // Return the default model if the provided model is not in the list - return modelList[0] + return BEDROCK_MODELS[0] } // Configure configures the AmazonBedRockClient with the provided configuration and language.