feat: add token usage tracking and configuration options

- Updated `wecoding-sdk-go` dependency from version `0.8.8` to `0.8.10`.
- Removed unused `json` import and simplified completion output handling.
- Added `Usage` field to `CompletionOutput` and `StreamCompletionOutput` types.
- Introduced `ShowTokenUsages` configuration option to display token usage in responses.
- Added `TokenUsage` field to the `Chat` struct to track token usage.
- Refactored `saveConversation` method to include token usage details in the output.
- Added markdown rendering for conversation save messages.

Signed-off-by: codiing-hui <wecoding@yeah.net>

Author: coding-hui

go.mod

@@ -21,7 +21,7 @@ require (
 	github.com/charmbracelet/x/exp/strings v0.0.0-20240722160745-212f7b056ed0
 	github.com/coding-hui/common v0.8.7
 	github.com/coding-hui/go-prompt v0.2.8
-	github.com/coding-hui/wecoding-sdk-go v0.8.8
+	github.com/coding-hui/wecoding-sdk-go v0.8.10
 	github.com/erikgeiser/promptkit v0.9.0
 	github.com/fatih/color v1.18.0
 	github.com/ghodss/yaml v1.0.0

go.sum

@@ -81,8 +81,8 @@ github.com/coding-hui/common v0.8.7 h1:f9iHZcdQLgRFW/nIJHVfBINuzzCKsUcwVIhfboxXe
 github.com/coding-hui/common v0.8.7/go.mod h1:KaNHOP9IYLgHzouuvR8/IijT4tdFyKMzB5ghI+ShtIM=
 github.com/coding-hui/go-prompt v0.2.8 h1:SjoClQQmsRxytxn2o3OaUB4V/PNOk8NPjPHpqNss1GQ=
 github.com/coding-hui/go-prompt v0.2.8/go.mod h1:ZZeEnbY5BcUJ7xBpRM2iPp8GJr5+lxRU4l0Tcd8bvW0=
-github.com/coding-hui/wecoding-sdk-go v0.8.8 h1:fSR4c+05tfn0bfbZG2SKKxTcocBGirnTWF6K96nCCKk=
-github.com/coding-hui/wecoding-sdk-go v0.8.8/go.mod h1:nWank9mlVeuHG2UEhHo/HwGr0p/ZnH4/3a5mjpLOOBw=
+github.com/coding-hui/wecoding-sdk-go v0.8.10 h1:u3TrDFJLthnC9IIviUSRc5uRKst1fw9LRXLmZ/gVEDg=
+github.com/coding-hui/wecoding-sdk-go v0.8.10/go.mod h1:nWank9mlVeuHG2UEhHo/HwGr0p/ZnH4/3a5mjpLOOBw=
 github.com/cpuguy83/go-md2man/v2 v2.0.4/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
 github.com/creack/pty v1.1.18 h1:n56/Zwd5o6whRC5PMGretI4IdRLlmBXYNjScPaBgsbY=
 github.com/creack/pty v1.1.18/go.mod h1:MOBLtS5ELjhRRrroQr9kyvTxUAFNvYEK993ew/Vr4O4=

internal/ai/ai.go

@@ -2,7 +2,6 @@ package ai
 
 import (
 	"context"
-	"encoding/json"
 	"fmt"
 	"strings"
 
@@ -77,17 +76,14 @@ func (e *Engine) CreateCompletion(ctx context.Context, messages []llms.ChatMessa
 
 	e.appendAssistantMessage(content)
 
-	var output CompletionOutput
-	err = json.Unmarshal([]byte(content), &output)
-	if err != nil {
-		output = CompletionOutput{
-			Command:     "",
-			Explanation: content,
-			Executable:  false,
-		}
-	}
+	e.running = false
 
-	return &output, nil
+	return &CompletionOutput{
+		Command:     "",
+		Explanation: content,
+		Executable:  false,
+		Usage:       rsp.Usage,
+	}, nil
 }
 
 func (e *Engine) CreateStreamCompletion(ctx context.Context, messages []llms.ChatMessage) (*StreamCompletionOutput, error) {
@@ -135,6 +131,7 @@ func (e *Engine) CreateStreamCompletion(ctx context.Context, messages []llms.Cha
 			Content:    "",
 			Last:       true,
 			Executable: executable,
+			Usage:      rsp.Usage,
 		}
 	}
 	e.running = false
@@ -145,6 +142,7 @@ func (e *Engine) CreateStreamCompletion(ctx context.Context, messages []llms.Cha
 		Content:    output,
 		Last:       true,
 		Executable: executable,
+		Usage:      rsp.Usage,
 	}, nil
 }
 

internal/ai/types.go

@@ -21,6 +21,8 @@ type CompletionOutput struct {
 	Command     string `json:"cmd"`
 	Explanation string `json:"exp"`
 	Executable  bool   `json:"exec"`
+
+	Usage llms.Usage `json:"usage"`
 }
 
 func (c CompletionOutput) GetCommand() string {
@@ -46,6 +48,8 @@ type StreamCompletionOutput struct {
 	Last       bool
 	Interrupt  bool
 	Executable bool
+
+	Usage llms.Usage `json:"usage"`
 }
 
 func (c StreamCompletionOutput) GetContent() string {
@@ -63,3 +67,7 @@ func (c StreamCompletionOutput) IsInterrupt() bool {
 func (c StreamCompletionOutput) IsExecutable() bool {
 	return c.Executable
 }
+
+func (c StreamCompletionOutput) GetUsage() llms.Usage {
+	return c.Usage
+}

internal/options/config.go

@@ -80,34 +80,36 @@ var Help = map[string]string{
 	"datastore":           "Configure the datastore to use.",
 	"auto-coder":          "Configure the auto coder to use.",
 	"auto-commit":         "Automatically commit code changes after generation.",
+	"show-token-usage":    "Show token usage in the response.",
 	"verbose":             "Verbose mode. 0: no verbose, 1: debug verbose",
 }
 
 // Config is a structure used to configure a AI.
 // Its members are sorted roughly in order of importance for composers.
 type Config struct {
-	Model         string     `yaml:"default-model" env:"MODEL"`
-	API           string     `yaml:"default-api" env:"API"`
-	Raw           bool       `yaml:"raw" env:"RAW"`
-	Quiet         bool       `yaml:"quiet" env:"QUIET"`
-	MaxTokens     int        `yaml:"max-tokens" env:"MAX_TOKENS"`
-	MaxInputChars int        `yaml:"max-input-chars" env:"MAX_INPUT_CHARS"`
-	Temperature   float64    `yaml:"temp" env:"TEMP"`
-	Stop          []string   `yaml:"stop" env:"STOP"`
-	TopP          float64    `yaml:"topp" env:"TOPP"`
-	TopK          int        `yaml:"topk" env:"TOPK"`
-	NoLimit       bool       `yaml:"no-limit" env:"NO_LIMIT"`
-	NoCache       bool       `yaml:"no-cache" env:"NO_CACHE"`
-	MaxRetries    int        `yaml:"max-retries" env:"MAX_RETRIES"`
-	WordWrap      int        `yaml:"word-wrap" env:"WORD_WRAP"`
-	Fanciness     uint       `yaml:"fanciness" env:"FANCINESS"`
-	LoadingText   string     `yaml:"loading-text" env:"LOADING_TEXT"`
-	FormatText    FormatText `yaml:"format-text"`
-	FormatAs      string     `yaml:"format-as" env:"FORMAT_AS"`
-	Verbose       int        `yaml:"verbose" env:"VERBOSE"`
-	APIs          APIs       `yaml:"apis"`
-	DataStore     DataStore  `yaml:"datastore"`
-	AutoCoder     AutoCoder  `yaml:"auto-coder"`
+	Model           string     `yaml:"default-model" env:"MODEL"`
+	API             string     `yaml:"default-api" env:"API"`
+	Raw             bool       `yaml:"raw" env:"RAW"`
+	Quiet           bool       `yaml:"quiet" env:"QUIET"`
+	MaxTokens       int        `yaml:"max-tokens" env:"MAX_TOKENS"`
+	MaxInputChars   int        `yaml:"max-input-chars" env:"MAX_INPUT_CHARS"`
+	Temperature     float64    `yaml:"temp" env:"TEMP"`
+	Stop            []string   `yaml:"stop" env:"STOP"`
+	TopP            float64    `yaml:"topp" env:"TOPP"`
+	TopK            int        `yaml:"topk" env:"TOPK"`
+	NoLimit         bool       `yaml:"no-limit" env:"NO_LIMIT"`
+	NoCache         bool       `yaml:"no-cache" env:"NO_CACHE"`
+	MaxRetries      int        `yaml:"max-retries" env:"MAX_RETRIES"`
+	WordWrap        int        `yaml:"word-wrap" env:"WORD_WRAP"`
+	Fanciness       uint       `yaml:"fanciness" env:"FANCINESS"`
+	LoadingText     string     `yaml:"loading-text" env:"LOADING_TEXT"`
+	FormatText      FormatText `yaml:"format-text"`
+	FormatAs        string     `yaml:"format-as" env:"FORMAT_AS"`
+	Verbose         int        `yaml:"verbose" env:"VERBOSE"`
+	APIs            APIs       `yaml:"apis"`
+	DataStore       DataStore  `yaml:"datastore"`
+	AutoCoder       AutoCoder  `yaml:"auto-coder"`
+	ShowTokenUsages bool       `yaml:"show-token-usage" env:"SHOW_TOKEN_USAGES"`
 
 	DefaultPromptMode string `yaml:"default-prompt-mode,omitempty"`
 	ConversationID    string `yaml:"convo-id,omitempty"`

internal/options/config_template.yml

@@ -47,6 +47,8 @@ loading-text: Generating
 theme: charm
 # {{ index .Help "max-input-chars" }}
 max-input-chars: 12250
+# {{ index .Help "show-token-usage" }}
+show-token-usage: true
 # {{ index .Help "max-tokens" }}
 # max-tokens: 100
 # {{ index .Help "datastore" }}

internal/ui/chat/chat.go

@@ -38,6 +38,7 @@ type Chat struct {
 	Error      *errbook.AiError
 	Output     string
 	GlamOutput string
+	TokenUsage llms.Usage
 
 	state  state
 	opts   *Options
@@ -82,7 +83,7 @@ func NewChat(cfg *options.Config, opts ...Option) *Chat {
 	}
 }
 
-// Run starts the chat.
+// Run starts the c.
 func (c *Chat) Run() error {
 	if _, err := tea.NewProgram(c).Run(); err != nil {
 		return errbook.Wrap("Couldn't start Bubble Tea program.", err)
@@ -110,7 +111,7 @@ func (c *Chat) Run() error {
 	}
 
 	if c.config.CacheWriteToID != "" {
-		return saveConversation(c)
+		return c.saveConversation()
 	}
 
 	return nil
@@ -158,6 +159,7 @@ func (c *Chat) Update(msg tea.Msg) (tea.Model, tea.Cmd) {
 		}
 		if msg.IsLast() {
 			c.state = doneState
+			c.TokenUsage = msg.GetUsage()
 			return c, c.quit
 		}
 		cmds = append(cmds, c.awaitChatCompletedCmd())
@@ -330,9 +332,20 @@ func (c *Chat) readStdinCmd() tea.Msg {
 	}
 }
 
-func saveConversation(chat *Chat) error {
-	if chat.config.NoCache {
-		if !chat.config.Quiet {
+func (c *Chat) renderMarkdown(raw string) string {
+	if c.config.Raw {
+		return raw
+	}
+	rendered, err := c.glam.Render(raw)
+	if err != nil {
+		return raw
+	}
+	return rendered
+}
+
+func (c *Chat) saveConversation() error {
+	if c.config.NoCache {
+		if !c.config.Quiet {
 			fmt.Fprintf(
 				os.Stderr,
 				"\nConversation was not saved because %s or %s is set.\n",
@@ -344,9 +357,9 @@ func saveConversation(chat *Chat) error {
 	}
 
 	ctx := context.Background()
-	convoStore := chat.engine.GetConvoStore()
-	writeToID := chat.config.CacheWriteToID
-	writeToTitle := strings.TrimSpace(chat.config.CacheWriteToTitle)
+	convoStore := c.engine.GetConvoStore()
+	writeToID := c.config.CacheWriteToID
+	writeToTitle := strings.TrimSpace(c.config.CacheWriteToTitle)
 
 	if convo.MatchSha1(writeToTitle) || writeToTitle == "" {
 		messages, err := convoStore.Messages(ctx, writeToID)
@@ -363,28 +376,32 @@ func saveConversation(chat *Chat) error {
 	if err := convoStore.PersistentMessages(ctx, writeToID); err != nil {
 		return errbook.Wrap(fmt.Sprintf(
 			"There was a problem writing %s to the cache. Use %s / %s to disable it.",
-			chat.config.CacheWriteToID,
+			c.config.CacheWriteToID,
 			console.StderrStyles().InlineCode.Render("--no-cache"),
 			console.StderrStyles().InlineCode.Render("NO_CACHE"),
 		), err)
 	}
 
-	if err := convoStore.SaveConversation(ctx, writeToID, writeToTitle, chat.config.Model); err != nil {
+	if err := convoStore.SaveConversation(ctx, writeToID, writeToTitle, c.config.Model); err != nil {
 		return errbook.Wrap(fmt.Sprintf(
 			"There was a problem writing %s to the cache. Use %s / %s to disable it.",
-			chat.config.CacheWriteToID,
+			c.config.CacheWriteToID,
 			console.StderrStyles().InlineCode.Render("--no-cache"),
 			console.StderrStyles().InlineCode.Render("NO_CACHE"),
 		), err)
 	}
 
-	if !chat.config.Quiet {
-		fmt.Fprintln(
-			os.Stderr,
-			"\nConversation saved:",
-			console.StderrStyles().InlineCode.Render(chat.config.CacheWriteToID[:convo.Sha1short]),
-			console.StderrStyles().Comment.Render(writeToTitle),
-		)
+	if !c.config.Quiet {
+		content := fmt.Sprintf("\n**Conversation successfully saved:** `%s` `%s`\n", c.config.CacheWriteToID[:convo.Sha1short], writeToTitle)
+		if c.config.ShowTokenUsages {
+			content += fmt.Sprintf("\nFirst Token: `%.3fs` | Avg: `%.3f/s` | Total: `%.3fs` | Tokens: `%d`",
+				c.TokenUsage.FirstTokenTime.Seconds(),
+				c.TokenUsage.AverageTokensPerSecond,
+				c.TokenUsage.TotalTime.Seconds(),
+				c.TokenUsage.TotalTokens,
+			)
+		}
+		_, _ = fmt.Fprint(os.Stderr, c.renderMarkdown(content))
 	}
 
 	return nil