fix bug

bricks-cloud · Jan 5, 2025 · 5fb4f93 · 5fb4f93
1 parent 34ebf34
commit 5fb4f93
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/internal/provider/anthropic/cost.go b/internal/provider/anthropic/cost.go
@@ -64,7 +64,13 @@ func (ce *CostEstimator) EstimatePromptCost(model string, tks int) (float64, err
 
 	}
 
-	selected := selectModel(model)
+	selected := ""
+	if strings.HasPrefix(model, "us") {
+		selected = convertAmazonModelToAnthropicModel(model)
+	} else {
+		selected = selectModel(model)
+	}
+
 	cost, ok := costMap[selected]
 	if !ok {
 		return 0, fmt.Errorf("%s is not present in the cost map provided", model)