From a4e53cb1fc0cc716d8a2287655bf6301142050bd Mon Sep 17 00:00:00 2001 From: northes Date: Thu, 6 Jun 2024 16:57:06 +0800 Subject: [PATCH] feat: bind usage in stream mode --- api_chat_completions.go | 6 +++++- api_chat_completions_test.go | 8 +++++++- 2 files changed, 12 insertions(+), 2 deletions(-) diff --git a/api_chat_completions.go b/api_chat_completions.go index e4dfe6e..21a3944 100644 --- a/api_chat_completions.go +++ b/api_chat_completions.go @@ -53,7 +53,8 @@ type ChatCompletionsResponse struct { Created int `json:"created"` Model string `json:"model"` Choices []*ChatCompletionsResponseChoices `json:"choices"` - Usage *ChatCompletionsResponseUsage `json:"usage"` + // returns only in non-stream mode + Usage *ChatCompletionsResponseUsage `json:"usage,omitempty"` } type ChatCompletionsResponseChoices struct { @@ -65,6 +66,9 @@ type ChatCompletionsResponseChoices struct { Delta *ChatCompletionsMessage `json:"delta,omitempty"` FinishReason ChatCompletionsFinishReason `json:"finish_reason"` + + // returns only in stream mode + Usage *ChatCompletionsResponseUsage `json:"usage,omitempty"` } type ChatCompletionsResponseUsage struct { diff --git a/api_chat_completions_test.go b/api_chat_completions_test.go index 378992c..ec2e8c0 100644 --- a/api_chat_completions_test.go +++ b/api_chat_completions_test.go @@ -55,6 +55,13 @@ func TestChatStream(t *testing.T) { msg, err := receive.GetMessage() if err != nil { if errors.Is(err, io.EOF) { + // Finish usage + if len(receive.Choices) != 0 { + choice := receive.Choices[0] + if choice.FinishReason == moonshot.FinishReasonStop && choice.Usage != nil { + t.Logf("Finish Usage: PromptTokens: %d, CompletionTokens: %d, TotalTokens: %d", choice.Usage.PromptTokens, choice.Usage.CompletionTokens, choice.Usage.TotalTokens) + } + } break } t.Error(err) @@ -68,7 +75,6 @@ func TestChatStream(t *testing.T) { default: t.Logf("Role: %s,Content: %s", msg.Role, msg.Content) } - } }