feat(run/chat): display token usage

Signed-off-by: Dorin Geman <[email protected]>

Author: doringeman

commands/run.go

@@ -271,7 +271,7 @@ func chatWithMarkdown(cmd *cobra.Command, client *desktop.Client, backend, model
 		// Simple case: just stream as plain text
 		return client.Chat(backend, model, prompt, apiKey, func(content string) {
 			cmd.Print(content)
-		})
+		}, false)
 	}
 
 	// For markdown: use streaming buffer to render code blocks as they complete
@@ -289,7 +289,7 @@ func chatWithMarkdown(cmd *cobra.Command, client *desktop.Client, backend, model
 		} else if rendered != "" {
 			cmd.Print(rendered)
 		}
-	})
+	}, true)
 	if err != nil {
 		return err
 	}

desktop/api.go

@@ -40,4 +40,9 @@ type OpenAIChatResponse struct {
 		Index        int    `json:"index"`
 		FinishReason string `json:"finish_reason"`
 	} `json:"choices"`
+	Usage *struct {
+		CompletionTokens int `json:"completion_tokens"`
+		PromptTokens     int `json:"prompt_tokens"`
+		TotalTokens      int `json:"total_tokens"`
+	} `json:"usage,omitempty"`
 }

desktop/desktop.go

@@ -365,7 +365,7 @@ func (c *Client) fullModelID(id string) (string, error) {
 }
 
 // Chat performs a chat request and streams the response content with selective markdown rendering.
-func (c *Client) Chat(backend, model, prompt, apiKey string, outputFunc func(string)) error {
+func (c *Client) Chat(backend, model, prompt, apiKey string, outputFunc func(string), shouldUseMarkdown bool) error {
 	model = normalizeHuggingFaceModelName(model)
 	if !strings.Contains(strings.Trim(model, "/"), "/") {
 		// Do an extra API call to check if the model parameter isn't a model ID.
@@ -422,7 +422,14 @@ func (c *Client) Chat(backend, model, prompt, apiKey string, outputFunc func(str
 	)
 
 	printerState := chatPrinterNone
-	reasoningFmt := color.New(color.FgWhite).Add(color.Italic)
+	reasoningFmt := color.New().Add(color.Italic)
+
+	var finalUsage *struct {
+		CompletionTokens int `json:"completion_tokens"`
+		PromptTokens     int `json:"prompt_tokens"`
+		TotalTokens      int `json:"total_tokens"`
+	}
+
 	scanner := bufio.NewScanner(resp.Body)
 	for scanner.Scan() {
 		line := scanner.Text()
@@ -445,6 +452,10 @@ func (c *Client) Chat(backend, model, prompt, apiKey string, outputFunc func(str
 			return fmt.Errorf("error parsing stream response: %w", err)
 		}
 
+		if streamResp.Usage != nil {
+			finalUsage = streamResp.Usage
+		}
+
 		if len(streamResp.Choices) > 0 {
 			if streamResp.Choices[0].Delta.ReasoningContent != "" {
 				chunk := streamResp.Choices[0].Delta.ReasoningContent
@@ -454,14 +465,14 @@ func (c *Client) Chat(backend, model, prompt, apiKey string, outputFunc func(str
 				if printerState != chatPrinterReasoning {
 					const thinkingHeader = "Thinking:\n"
 					if reasoningFmt != nil {
-						outputFunc(reasoningFmt.Sprint(thinkingHeader))
+						reasoningFmt.Print(thinkingHeader)
 					} else {
 						outputFunc(thinkingHeader)
 					}
 				}
 				printerState = chatPrinterReasoning
 				if reasoningFmt != nil {
-					outputFunc(reasoningFmt.Sprint(chunk))
+					reasoningFmt.Print(chunk)
 				} else {
 					outputFunc(chunk)
 				}
@@ -481,6 +492,19 @@ func (c *Client) Chat(backend, model, prompt, apiKey string, outputFunc func(str
 		return fmt.Errorf("error reading response stream: %w", err)
 	}
 
+	if finalUsage != nil {
+		usageInfo := fmt.Sprintf("\n\nToken usage: %d prompt + %d completion = %d total",
+			finalUsage.PromptTokens,
+			finalUsage.CompletionTokens,
+			finalUsage.TotalTokens)
+
+		usageFmt := color.New(color.FgHiBlack)
+		if !shouldUseMarkdown {
+			usageFmt.DisableColor()
+		}
+		outputFunc(usageFmt.Sprint(usageInfo))
+	}
+
 	return nil
 }
 

desktop/desktop_test.go

@@ -63,7 +63,7 @@ func TestChatHuggingFaceModel(t *testing.T) {
 		Body:       io.NopCloser(bytes.NewBufferString("data: {\"choices\":[{\"delta\":{\"content\":\"Hello there!\"}}]}\n")),
 	}, nil)
 
-	err := client.Chat("", modelName, prompt, "", func(s string) {})
+	err := client.Chat("", modelName, prompt, "", func(s string) {}, false)
 	assert.NoError(t, err)
 }