From d139f36efda9fb2870a5e5e27e66c0dce03ff64a Mon Sep 17 00:00:00 2001 From: Adrien Date: Wed, 22 May 2024 20:09:03 +0200 Subject: [PATCH] fix --- RequestAnthropic.go | 5 ++++- RequestGoogle.go | 5 ++++- RequestMistral.go | 6 +++++- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/RequestAnthropic.go b/RequestAnthropic.go index bc1d744..75ba6d0 100644 --- a/RequestAnthropic.go +++ b/RequestAnthropic.go @@ -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 = $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 diff --git a/RequestGoogle.go b/RequestGoogle.go index 78efb40..362e27e 100644 --- a/RequestGoogle.go +++ b/RequestGoogle.go @@ -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 = $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 diff --git a/RequestMistral.go b/RequestMistral.go index 92f4147..895efca 100644 --- a/RequestMistral.go +++ b/RequestMistral.go @@ -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)