This commit is contained in:
Adrien Bouvais 2024-05-22 20:09:03 +02:00
parent d5ab5a8e7c
commit d139f36efd
3 changed files with 13 additions and 3 deletions

View File

@ -164,7 +164,7 @@ func RequestAnthropic(model string, messages []Message, maxTokens int, temperatu
}
var usedModelInfo ModelInfo
edgeClient.QuerySingle(edgeCtx, `
err = edgeClient.QuerySingle(edgeCtx, `
SELECT ModelInfo {
inputPrice,
outputPrice
@ -172,6 +172,9 @@ func RequestAnthropic(model string, messages []Message, maxTokens int, temperatu
FILTER .modelID = <str>$0
LIMIT 1
`, &usedModelInfo, model)
if err != nil {
return AnthropicChatCompletionResponse{}, fmt.Errorf("error getting model info: %w", err)
}
var inputCost float32 = float32(chatCompletionResponse.Usage.InputTokens) * usedModelInfo.InputPrice
var outputCost float32 = float32(chatCompletionResponse.Usage.OutputTokens) * usedModelInfo.OutputPrice

View File

@ -171,7 +171,7 @@ func RequestGoogle(model string, messages []Message, temperature float64) (Opena
}
var usedModelInfo ModelInfo
edgeClient.QuerySingle(edgeCtx, `
err = edgeClient.QuerySingle(edgeCtx, `
SELECT ModelInfo {
inputPrice,
outputPrice
@ -179,6 +179,9 @@ func RequestGoogle(model string, messages []Message, temperature float64) (Opena
FILTER .modelID = <str>$0
LIMIT 1
`, &usedModelInfo, model)
if err != nil {
return OpenaiChatCompletionResponse{}, fmt.Errorf("error getting model info: %w", err)
}
var inputCost float32 = float32(chatCompletionResponse.Usage.PromptTokens) * usedModelInfo.InputPrice
var outputCost float32 = float32(chatCompletionResponse.Usage.CompletionTokens) * usedModelInfo.OutputPrice

View File

@ -169,7 +169,7 @@ func RequestMistral(model string, messages []Message, temperature float64) (Mist
}
var usedModelInfo ModelInfo
edgeClient.QuerySingle(edgeCtx, `
err = edgeClient.QuerySingle(edgeCtx, `
SELECT ModelInfo {
inputPrice,
outputPrice
@ -178,6 +178,10 @@ func RequestMistral(model string, messages []Message, temperature float64) (Mist
LIMIT 1
`, &usedModelInfo, model)
if usedModelInfo.InputPrice == 0 || usedModelInfo.OutputPrice == 0 {
return MistralChatCompletionResponse{}, fmt.Errorf("model %s not found in Mistral", model)
}
var inputCost float32 = float32(chatCompletionResponse.Usage.PromptTokens) * usedModelInfo.InputPrice
var outputCost float32 = float32(chatCompletionResponse.Usage.CompletionTokens) * usedModelInfo.OutputPrice
addUsage(inputCost, outputCost, chatCompletionResponse.Usage.PromptTokens, chatCompletionResponse.Usage.CompletionTokens, model)