stream: make usage optional on streamed response until final chunk

joinhive-ai · Oct 30, 2024 · a2777bc · a2777bc
1 parent 253a786
commit a2777bc
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/common/utils.go b/common/utils.go
@@ -92,7 +92,7 @@ func TransformResponse(req *worker.GenLLMFormdataRequestBody, resp *http.Respons
 				FinishReason: "stop",
 			},
 		},
-		Usage: models.Usage{
+		Usage: &models.Usage{
 			TotalTokens: res.TokensUsed, // TokensUsed already includes prompt tokens
 		},
 	}
@@ -120,7 +120,7 @@ func TransformStreamResponse(chunk worker.LlmStreamChunk, req *worker.GenLLMForm
 	if chunk.Done {
 		openAIResp.Choices[0].FinishReason = "stop"
 		// Only include usage information in the final chunk
-		openAIResp.Usage = models.Usage{
+		openAIResp.Usage = &models.Usage{
 			TotalTokens: chunk.TokensUsed, // TokensUsed already includes prompt tokens
 		}
 	}

diff --git a/models/models.go b/models/models.go
@@ -25,7 +25,7 @@ type OpenAIResponse struct {
 	Created int64    `json:"created"`
 	Model   string   `json:"model"`
 	Choices []Choice `json:"choices"`
-	Usage   Usage    `json:"usage"`
+	Usage   *Usage   `json:"usage"`
 }
 
 type OpenAIStreamResponse struct {
@@ -34,7 +34,7 @@ type OpenAIStreamResponse struct {
 	Created int64          `json:"created"`
 	Model   string         `json:"model"`
 	Choices []StreamChoice `json:"choices"`
-	Usage   Usage          `json:"usage"`
+	Usage   *Usage         `json:"usage,omitempty"`
 }
 
 type StreamChoice struct {