Skip to content

Commit

Permalink
add support for amazon bedrock model
Browse files Browse the repository at this point in the history
  • Loading branch information
spikelu2016 committed Dec 27, 2024
1 parent c101ce8 commit d88f94b
Showing 1 changed file with 16 additions and 1 deletion.
17 changes: 16 additions & 1 deletion internal/provider/anthropic/cost.go
Original file line number Diff line number Diff line change
Expand Up @@ -94,13 +94,28 @@ func selectModel(model string) string {
return ""
}

func convertAmazonModelToAnthropicModel(model string) string {
parts := strings.Split(model, ".")
if len(parts) < 3 {
return model
}

return selectModel(parts[2])
}

func (ce *CostEstimator) EstimateCompletionCost(model string, tks int) (float64, error) {
costMap, ok := ce.tokenCostMap["completion"]
if !ok {
return 0, errors.New("prompt token cost is not provided")
}

selected := selectModel(model)
selected := ""
if strings.HasPrefix(model, "us") {
selected = convertAmazonModelToAnthropicModel(model)
} else {
selected = selectModel(model)
}

cost, ok := costMap[selected]
if !ok {
return 0, errors.New("model is not present in the cost map provided")
Expand Down

0 comments on commit d88f94b

Please sign in to comment.