feat: bind usage in stream mode

northes · Jun 6, 2024 · a4e53cb · a4e53cb
1 parent 5bfd815
commit a4e53cb
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 2 deletions.
diff --git 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
@@ -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)
 		}
-
 	}
 }