Merge pull request #18 from northes/stream_mode_usage

feat: bind usage info in stream mode
northes · Jun 6, 2024 · 5d9275b · 5d9275b
2 parents 85f703d + a4e53cb
commit 5d9275b
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 3 deletions.
diff --git a/README.md b/README.md
@@ -75,7 +75,7 @@ if !ok {
 
 cli, err := moonshot.NewClientWithConfig(
     moonshot.NewConfig(
-        moonshot.WithAPIKey("xxxx"),
+        moonshot.WithAPIKey(key),
     ),
 )
 ```

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)
 		}
-
 	}
 }