chore: collect and log token usage

Author: ibuildthecloud

go.mod

@@ -14,7 +14,7 @@ require (
 	github.com/fatih/color v1.16.0
 	github.com/getkin/kin-openapi v0.123.0
 	github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510
-	github.com/gptscript-ai/chat-completion-client v0.0.0-20240502162133-7dabc28eab59
+	github.com/gptscript-ai/chat-completion-client v0.0.0-20240515050533-bdef9f2226a9
 	github.com/hexops/autogold/v2 v2.2.1
 	github.com/jaytaylor/html2text v0.0.0-20230321000545-74c2419ad056
 	github.com/mholt/archiver/v4 v4.0.0-alpha.8

go.sum

@@ -123,8 +123,8 @@ github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+
 github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=
 github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc=
 github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
-github.com/gptscript-ai/chat-completion-client v0.0.0-20240502162133-7dabc28eab59 h1:Nda0GDkrmIDiAFHfaXu0eIp6SsBs0/4Zyo87ff3Qcqo=
-github.com/gptscript-ai/chat-completion-client v0.0.0-20240502162133-7dabc28eab59/go.mod h1:7P/o6/IWa1KqsntVf68hSnLKuu3+xuqm6lYhch1w4jo=
+github.com/gptscript-ai/chat-completion-client v0.0.0-20240515050533-bdef9f2226a9 h1:s6nL/aokB1sJTqVXEjN0zFI5CJa66ubw9g68VTMzEw0=
+github.com/gptscript-ai/chat-completion-client v0.0.0-20240515050533-bdef9f2226a9/go.mod h1:7P/o6/IWa1KqsntVf68hSnLKuu3+xuqm6lYhch1w4jo=
 github.com/hashicorp/errwrap v1.0.0 h1:hLrqtEDnRye3+sgx6z4qVLNuviH3MR5aQ0ykNJa/UYA=
 github.com/hashicorp/errwrap v1.0.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4=
 github.com/hashicorp/go-multierror v1.1.1 h1:H5DkEtf6CXdFp0N0Em5UCwQpXMWke8IA0+lD48awMYo=

pkg/monitor/display.go

@@ -62,6 +62,7 @@ type display struct {
 	dumpState     string
 	callIDMap     map[string]string
 	callLock      sync.Mutex
+	usage         types.Usage
 }
 
 type livePrinter struct {
@@ -226,6 +227,10 @@ func (d *display) Event(event runner.Event) {
 		userSpecifiedToolName: event.CallContext.ToolName,
 	}
 
+	d.usage.PromptTokens += event.Usage.PromptTokens
+	d.usage.CompletionTokens += event.Usage.CompletionTokens
+	d.usage.TotalTokens += event.Usage.TotalTokens
+
 	switch event.Type {
 	case runner.EventTypeCallStart:
 		d.livePrinter.progressStart(currentCall)
@@ -283,6 +288,7 @@ func (d *display) Stop(output string, err error) {
 	defer d.callLock.Unlock()
 
 	log.Fields("runID", d.dump.ID, "output", output, "err", err).Debugf("Run stopped")
+	log.Fields("runID", d.dump.ID, "total", d.usage.TotalTokens, "prompt", d.usage.PromptTokens, "completion", d.usage.CompletionTokens).Infof("usage   ")
 	d.dump.Output = output
 	d.dump.Err = err
 	if d.dumpState != "" {

pkg/mvl/log.go

@@ -38,6 +38,15 @@ func (f formatter) Format(entry *logrus.Entry) ([]byte, error) {
 	if i, ok := entry.Data["response"]; ok && i != "" {
 		msg += fmt.Sprintf(" [response=%s]", i)
 	}
+	if i, ok := entry.Data["total"]; ok && i != "" {
+		msg += fmt.Sprintf(" [total=%v]", i)
+	}
+	if i, ok := entry.Data["prompt"]; ok && i != "" {
+		msg += fmt.Sprintf(" [prompt=%v]", i)
+	}
+	if i, ok := entry.Data["completion"]; ok && i != "" {
+		msg += fmt.Sprintf(" [completion=%v]", i)
+	}
 	return []byte(fmt.Sprintf("%s %s\n",
 		entry.Time.Format(time.TimeOnly),
 		msg)), nil

pkg/openai/client.go

@@ -308,6 +308,9 @@ func (c *Client) Call(ctx context.Context, messageRequest types.CompletionReques
 		Model:     messageRequest.Model,
 		Messages:  msgs,
 		MaxTokens: messageRequest.MaxTokens,
+		StreamOptions: &openai.StreamOptions{
+			IncludeUsage: true,
+		},
 	}
 
 	if messageRequest.Temperature == nil {
@@ -372,17 +375,27 @@ func (c *Client) Call(ctx context.Context, messageRequest types.CompletionReques
 		}
 	}
 
+	var usage types.Usage
+	if !cacheResponse {
+		usage = result.Usage
+	}
+
 	status <- types.CompletionStatus{
 		CompletionID: id,
 		Chunks:       response,
 		Response:     result,
+		Usage:        usage,
 		Cached:       cacheResponse,
 	}
 
 	return &result, nil
 }
 
 func appendMessage(msg types.CompletionMessage, response openai.ChatCompletionStreamResponse) types.CompletionMessage {
+	msg.Usage.CompletionTokens = types.FirstSet(msg.Usage.CompletionTokens, response.Usage.CompletionTokens)
+	msg.Usage.PromptTokens = types.FirstSet(msg.Usage.PromptTokens, response.Usage.PromptTokens)
+	msg.Usage.TotalTokens = types.FirstSet(msg.Usage.TotalTokens, response.Usage.TotalTokens)
+
 	if len(response.Choices) == 0 {
 		return msg
 	}
@@ -470,6 +483,7 @@ func (c *Client) call(ctx context.Context, request openai.ChatCompletionRequest,
 				Object:  resp.Object,
 				Created: resp.Created,
 				Model:   resp.Model,
+				Usage:   resp.Usage,
 				Choices: []openai.ChatCompletionStreamChoice{
 					{
 						Index: resp.Choices[0].Index,

pkg/runner/runner.go

@@ -196,6 +196,7 @@ type Event struct {
 	ChatCompletionID   string                 `json:"chatCompletionId,omitempty"`
 	ChatRequest        any                    `json:"chatRequest,omitempty"`
 	ChatResponse       any                    `json:"chatResponse,omitempty"`
+	Usage              types.Usage            `json:"usage,omitempty"`
 	ChatResponseCached bool                   `json:"chatResponseCached,omitempty"`
 	Content            string                 `json:"content,omitempty"`
 }
@@ -622,6 +623,7 @@ func streamProgress(callCtx *engine.Context, monitor Monitor) (chan<- types.Comp
 					ChatCompletionID:   status.CompletionID,
 					ChatRequest:        status.Request,
 					ChatResponse:       status.Response,
+					Usage:              status.Usage,
 					ChatResponseCached: status.Cached,
 				})
 			}

pkg/types/completion.go

@@ -54,12 +54,20 @@ type CompletionMessage struct {
 	// ToolCall should be set for only messages of type "tool" and Content[0].Text should be set as the
 	// result of the call describe by this field
 	ToolCall *CompletionToolCall `json:"toolCall,omitempty"`
+	Usage    Usage               `json:"usage,omitempty"`
+}
+
+type Usage struct {
+	PromptTokens     int `json:"promptTokens,omitempty"`
+	CompletionTokens int `json:"completionTokens,omitempty"`
+	TotalTokens      int `json:"totalTokens,omitempty"`
 }
 
 type CompletionStatus struct {
 	CompletionID    string
 	Request         any
 	Response        any
+	Usage           Usage
 	Cached          bool
 	Chunks          any
 	PartialResponse *CompletionMessage