feat: consolidate token usage and test coverage (#223)

- Replace manual token printing with a single usage string call
- Introduce a method that handles token counts and their details
- Add tests verifying the usage string output in various scenarios

Signed-off-by: Bo-Yi Wu <[email protected]>

Author: appleboy

cmd/commit.go

@@ -4,7 +4,6 @@ import (
 	"html"
 	"os"
 	"path"
-	"strconv"
 	"strings"
 	"time"
 
@@ -149,10 +148,7 @@ var commitCmd = &cobra.Command{
 				return err
 			}
 			data[prompt.SummarizeMessageKey] = strings.TrimSpace(resp.Content)
-			color.Magenta("PromptTokens: " + strconv.Itoa(resp.Usage.PromptTokens) +
-				", CompletionTokens: " + strconv.Itoa(resp.Usage.CompletionTokens) +
-				", TotalTokens: " + strconv.Itoa(resp.Usage.TotalTokens),
-			)
+			color.Magenta(resp.Usage.String())
 		}
 
 		// Get summarize title from diff datas
@@ -174,10 +170,7 @@ var commitCmd = &cobra.Command{
 				return err
 			}
 			summarizeTitle := resp.Content
-			color.Magenta("PromptTokens: " + strconv.Itoa(resp.Usage.PromptTokens) +
-				", CompletionTokens: " + strconv.Itoa(resp.Usage.CompletionTokens) +
-				", TotalTokens: " + strconv.Itoa(resp.Usage.TotalTokens),
-			)
+			color.Magenta(resp.Usage.String())
 
 			// lowercase the first character of first word of the commit message and remove last period
 			summarizeTitle = strings.TrimRight(strings.ToLower(string(summarizeTitle[0]))+summarizeTitle[1:], ".")
@@ -203,10 +196,7 @@ var commitCmd = &cobra.Command{
 			}
 			summaryPrix = resp.Content
 
-			color.Magenta("PromptTokens: " + strconv.Itoa(resp.Usage.PromptTokens) +
-				", CompletionTokens: " + strconv.Itoa(resp.Usage.CompletionTokens) +
-				", TotalTokens: " + strconv.Itoa(resp.Usage.TotalTokens),
-			)
+			color.Magenta(resp.Usage.String())
 
 			data[prompt.SummarizePrefixKey] = summaryPrix
 		}
@@ -260,10 +250,7 @@ var commitCmd = &cobra.Command{
 			if err != nil {
 				return err
 			}
-			color.Magenta("PromptTokens: " + strconv.Itoa(resp.Usage.PromptTokens) +
-				", CompletionTokens: " + strconv.Itoa(resp.Usage.CompletionTokens) +
-				", TotalTokens: " + strconv.Itoa(resp.Usage.TotalTokens),
-			)
+			color.Magenta(resp.Usage.String())
 			commitMessage = resp.Content
 		}
 

core/openai.go

@@ -2,6 +2,7 @@ package core
 
 import (
 	"context"
+	"strconv"
 
 	"github.com/sashabaranov/go-openai"
 )
@@ -18,6 +19,19 @@ type Usage struct {
 	CompletionTokensDetails *openai.CompletionTokensDetails
 }
 
+func (u Usage) String() string {
+	s := "Prompt tokens: " + strconv.Itoa(u.PromptTokens)
+	if u.PromptTokensDetails != nil && u.PromptTokensDetails.CachedTokens > 0 {
+		s += " (CachedTokens: " + strconv.Itoa(u.PromptTokensDetails.CachedTokens) + ")"
+	}
+	s += ", Completion tokens: " + strconv.Itoa(u.CompletionTokens)
+	if u.CompletionTokensDetails != nil && u.CompletionTokensDetails.ReasoningTokens > 0 {
+		s += " (ReasoningTokens: " + strconv.Itoa(u.CompletionTokensDetails.ReasoningTokens) + ")"
+	}
+	s += ", Total tokens: " + strconv.Itoa(u.TotalTokens)
+	return s
+}
+
 // Response represents the structure of a response from the OpenAI API.
 // It contains the content of the response and the usage information.
 type Response struct {

core/openai_test.go

@@ -0,0 +1,65 @@
+package core
+
+import (
+	"testing"
+
+	"github.com/sashabaranov/go-openai"
+)
+
+func TestUsageString(t *testing.T) {
+	tests := []struct {
+		name     string
+		usage    Usage
+		expected string
+	}{
+		{
+			name: "without details",
+			usage: Usage{
+				PromptTokens:     10,
+				CompletionTokens: 20,
+				TotalTokens:      30,
+			},
+			expected: "Prompt tokens: 10, Completion tokens: 20, Total tokens: 30",
+		},
+		{
+			name: "with cached tokens",
+			usage: Usage{
+				PromptTokens:        10,
+				CompletionTokens:    20,
+				TotalTokens:         30,
+				PromptTokensDetails: &openai.PromptTokensDetails{CachedTokens: 5},
+			},
+			expected: "Prompt tokens: 10 (CachedTokens: 5), Completion tokens: 20, Total tokens: 30",
+		},
+		{
+			name: "with reasoning tokens",
+			usage: Usage{
+				PromptTokens:            10,
+				CompletionTokens:        20,
+				TotalTokens:             30,
+				CompletionTokensDetails: &openai.CompletionTokensDetails{ReasoningTokens: 7},
+			},
+			expected: "Prompt tokens: 10, Completion tokens: 20 (ReasoningTokens: 7), Total tokens: 30",
+		},
+		{
+			name: "with both details",
+			usage: Usage{
+				PromptTokens:            15,
+				CompletionTokens:        25,
+				TotalTokens:             40,
+				PromptTokensDetails:     &openai.PromptTokensDetails{CachedTokens: 3},
+				CompletionTokensDetails: &openai.CompletionTokensDetails{ReasoningTokens: 9},
+			},
+			expected: "Prompt tokens: 15 (CachedTokens: 3), Completion tokens: 25 (ReasoningTokens: 9), Total tokens: 40",
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			result := tc.usage.String()
+			if result != tc.expected {
+				t.Errorf("Usage.String() = %q, expected %q", result, tc.expected)
+			}
+		})
+	}
+}