From da94ebfd86298845974746b047b89bf25a3436c7 Mon Sep 17 00:00:00 2001 From: Luca Ronca Date: Thu, 6 Feb 2025 18:11:08 +0100 Subject: [PATCH] fix(bedrock): use a more flexible solution to extract provider name from the model id --- .../internal/bedrockclient/bedrockclient.go | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/llms/bedrock/internal/bedrockclient/bedrockclient.go b/llms/bedrock/internal/bedrockclient/bedrockclient.go index 2b3d2a313..e4f49c179 100644 --- a/llms/bedrock/internal/bedrockclient/bedrockclient.go +++ b/llms/bedrock/internal/bedrockclient/bedrockclient.go @@ -29,7 +29,20 @@ type Message struct { } func getProvider(modelID string) string { - return strings.Split(modelID, ".")[0] + switch { + case strings.Contains(modelID, "ai21"): + return "ai21" + case strings.Contains(modelID, "amazon"): + return "amazon" + case strings.Contains(modelID, "anthropic"): + return "anthropic" + case strings.Contains(modelID, "cohere"): + return "cohere" + case strings.Contains(modelID, "meta"): + return "meta" + default: + return "" + } } // NewClient creates a new Bedrock client.