diff --git a/cmd/cmd.go b/cmd/cmd.go index 0cafb4c1b..a271dc0bb 100644 --- a/cmd/cmd.go +++ b/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 { diff --git a/cmd/interactive.go b/cmd/interactive.go index a3c98c68e..e36ed800d 100644 --- a/cmd/interactive.go +++ b/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 Load a session or model") fmt.Fprintln(os.Stderr, " /save 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) }