Use a more flexible solution to extract provider name from the model id

tmc · Feb 6, 2025 · 2dfcb70 · 2dfcb70
1 parent 0672790
commit 2dfcb70
Showing 1 changed file with 14 additions and 1 deletion.
diff --git 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.