Ecosystem/ai-gateway
5
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

🐛 fix: Cohere did not calculate usage in stream mode.

Browse Source
This commit is contained in:
MartialBE 2024-05-11 15:01:02 +08:00
parent bc11c65c78
commit 5e2954a1a0
No known key found for this signature in database
GPG Key ID: 27C0267EC84B0A5C
1 changed files with 13 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
providers/cohere/chat.go
View File

@ -144,13 +144,9 @@ func ConvertToChatOpenai(provider base.ProviderInterface, response *CohereRespon
Created: common.GetTimestamp(), Created: common.GetTimestamp(),
Choices: []types.ChatCompletionChoice{choice}, Choices: []types.ChatCompletionChoice{choice},
Model: request.Model, Model: request.Model,
Usage: &types.Usage{ Usage: &types.Usage{},
PromptTokens: response.Meta.BilledUnits.InputTokens,
},
} }
*openaiResponse.Usage = usageHandle(&response.Meta.BilledUnits)
openaiResponse.Usage.CompletionTokens = response.Meta.BilledUnits.OutputTokens + response.Meta.Tokens.SearchUnits + response.Meta.Tokens.Classifications
openaiResponse.Usage.TotalTokens = openaiResponse.Usage.PromptTokens + openaiResponse.Usage.CompletionTokens
usage := provider.GetUsage() usage := provider.GetUsage()
*usage = *openaiResponse.Usage *usage = *openaiResponse.Usage
@ -188,6 +184,7 @@ func (h *CohereStreamHandler) convertToOpenaiStream(cohereResponse *CohereStream
if cohereResponse.EventType == "stream-end" { if cohereResponse.EventType == "stream-end" {
choice.FinishReason = types.FinishReasonStop choice.FinishReason = types.FinishReasonStop
*h.Usage = usageHandle(&cohereResponse.Response.Meta.BilledUnits)
} else { } else {
choice.Delta = types.ChatCompletionStreamChoiceDelta{ choice.Delta = types.ChatCompletionStreamChoiceDelta{
Role: types.ChatMessageRoleAssistant, Role: types.ChatMessageRoleAssistant,
@ -206,3 +203,13 @@ func (h *CohereStreamHandler) convertToOpenaiStream(cohereResponse *CohereStream
responseBody, _ := json.Marshal(chatCompletion) responseBody, _ := json.Marshal(chatCompletion)
dataChan <- string(responseBody) dataChan <- string(responseBody)
} }
func usageHandle(token *Tokens) types.Usage {
usage := types.Usage{
PromptTokens: token.InputTokens,
CompletionTokens: token.OutputTokens + token.SearchUnits + token.Classifications,
}
usage.TotalTokens = usage.PromptTokens + usage.CompletionTokens
return usage
}