Skip to content

Commit d88f94b

Browse files
committed
add support for amazon bedrock model
1 parent c101ce8 commit d88f94b

1 file changed

Lines changed: 16 additions & 1 deletion

File tree

internal/provider/anthropic/cost.go

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,13 +94,28 @@ func selectModel(model string) string {
9494
return ""
9595
}
9696

97+
func convertAmazonModelToAnthropicModel(model string) string {
98+
parts := strings.Split(model, ".")
99+
if len(parts) < 3 {
100+
return model
101+
}
102+
103+
return selectModel(parts[2])
104+
}
105+
97106
func (ce *CostEstimator) EstimateCompletionCost(model string, tks int) (float64, error) {
98107
costMap, ok := ce.tokenCostMap["completion"]
99108
if !ok {
100109
return 0, errors.New("prompt token cost is not provided")
101110
}
102111

103-
selected := selectModel(model)
112+
selected := ""
113+
if strings.HasPrefix(model, "us") {
114+
selected = convertAmazonModelToAnthropicModel(model)
115+
} else {
116+
selected = selectModel(model)
117+
}
118+
104119
cost, ok := costMap[selected]
105120
if !ok {
106121
return 0, errors.New("model is not present in the cost map provided")

0 commit comments

Comments
 (0)