cmd: add usage cmd to chat to see token consumption

Adding a `/usage` command to interactive cli chat sessions that displays the tokens used in the current sessions. This can be used alongside the models context window to understand when a context shift is going to happen.

cmd/cmd.go

@@ -1419,10 +1419,10 @@ func thinkingOutputClosingText(plainText bool) string {
 	return readline.ColorGrey + readline.ColorBold + text + readline.ColorDefault
 }
 
-func chat(cmd *cobra.Command, opts runOptions) (*api.Message, error) {
+func chat(cmd *cobra.Command, opts runOptions) (*api.Message, *api.Metrics, error) {
 	client, err := api.ClientFromEnvironment()
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	p := progress.NewProgress(os.Stderr)
@@ -1515,7 +1515,7 @@ func chat(cmd *cobra.Command, opts runOptions) (*api.Message, error) {
 
 	if err := client.Chat(cancelCtx, req, fn); err != nil {
 		if errors.Is(err, context.Canceled) {
-			return nil, nil
+			return nil, nil, nil
 		}
 
 		// this error should ideally be wrapped properly by the client
@@ -1523,9 +1523,9 @@ func chat(cmd *cobra.Command, opts runOptions) (*api.Message, error) {
 			p.StopAndClear()
 			fmt.Println("An error occurred while processing your message. Please try again.")
 			fmt.Println()
-			return nil, nil
+			return nil, nil, nil
 		}
-		return nil, err
+		return nil, nil, err
 	}
 
 	if len(opts.Messages) > 0 {
@@ -1535,14 +1535,14 @@ func chat(cmd *cobra.Command, opts runOptions) (*api.Message, error) {
 
 	verbose, err := cmd.Flags().GetBool("verbose")
 	if err != nil {
-		return nil, err
+		return nil, nil, err
 	}
 
 	if verbose {
 		latest.Summary()
 	}
 
-	return &api.Message{Role: role, Thinking: thinkingContent.String(), Content: fullResponse.String()}, nil
+	return &api.Message{Role: role, Thinking: thinkingContent.String(), Content: fullResponse.String()}, &latest.Metrics, nil
 }
 
 func generate(cmd *cobra.Command, opts runOptions) error {

cmd/interactive.go

@@ -30,6 +30,9 @@ const (
 )
 
 func generateInteractive(cmd *cobra.Command, opts runOptions) error {
+	var sessionPromptTokens int64
+	var sessionCompletionTokens int64
+
 	usage := func() {
 		fmt.Fprintln(os.Stderr, "Available Commands:")
 		fmt.Fprintln(os.Stderr, "  /set            Set session variables")
@@ -37,6 +40,7 @@ func generateInteractive(cmd *cobra.Command, opts runOptions) error {
 		fmt.Fprintln(os.Stderr, "  /load <model>   Load a session or model")
 		fmt.Fprintln(os.Stderr, "  /save <model>   Save your current session")
 		fmt.Fprintln(os.Stderr, "  /clear          Clear session context")
+		fmt.Fprintln(os.Stderr, "  /usage          Show session token usage")
 		fmt.Fprintln(os.Stderr, "  /bye            Exit")
 		fmt.Fprintln(os.Stderr, "  /?, /help       Help for a command")
 		fmt.Fprintln(os.Stderr, "  /? shortcuts    Help for keyboard shortcuts")
@@ -445,6 +449,9 @@ func generateInteractive(cmd *cobra.Command, opts runOptions) error {
 			} else {
 				usageShow()
 			}
+		case strings.HasPrefix(line, "/usage"):
+			fmt.Printf("prompt tokens:     %d\n", sessionPromptTokens)
+			fmt.Printf("completion tokens: %d\n", sessionCompletionTokens)
 		case strings.HasPrefix(line, "/help"), strings.HasPrefix(line, "/?"):
 			args := strings.Fields(line)
 			if len(args) > 1 {
@@ -499,7 +506,7 @@ func generateInteractive(cmd *cobra.Command, opts runOptions) error {
 
 			opts.Messages = append(opts.Messages, newMessage)
 
-			assistant, err := chat(cmd, opts)
+			assistant, metrics, err := chat(cmd, opts)
 			if err != nil {
 				if strings.Contains(err.Error(), "does not support thinking") ||
 					strings.Contains(err.Error(), "invalid think value") {
@@ -509,6 +516,10 @@ func generateInteractive(cmd *cobra.Command, opts runOptions) error {
 				}
 				return err
 			}
+			if metrics != nil {
+				sessionPromptTokens += int64(metrics.PromptEvalCount)
+				sessionCompletionTokens += int64(metrics.EvalCount)
+			}
 			if assistant != nil {
 				opts.Messages = append(opts.Messages, *assistant)
 			}