Skip to content

Commit 71261aa

Browse files
ibuildthecloudibuildthecloud
committed
chore: collect and log token usage
1 parent 12e80c2 commit 71261aa

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

49 files changed

+338
-151
lines changed

‎go.mod‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ require (
1414
github.com/fatih/color v1.16.0
1515
github.com/getkin/kin-openapi v0.123.0
1616
github.com/google/shlex v0.0.0-20191202100458-e7afc7fbc510
17-
github.com/gptscript-ai/chat-completion-client v0.0.0-20240502162133-7dabc28eab59
17+
github.com/gptscript-ai/chat-completion-client v0.0.0-20240515050533-bdef9f2226a9
1818
github.com/hexops/autogold/v2 v2.2.1
1919
github.com/jaytaylor/html2text v0.0.0-20230321000545-74c2419ad056
2020
github.com/mholt/archiver/v4 v4.0.0-alpha.8

‎go.sum‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -123,8 +123,8 @@ github.com/googleapis/gax-go/v2 v2.0.4/go.mod h1:0Wqv26UfaUD9n4G6kQubkQ+KchISgw+
123123
github.com/googleapis/gax-go/v2 v2.0.5/go.mod h1:DWXyrwAJ9X0FpwwEdw+IPEYBICEFu5mhpdKc/us6bOk=
124124
github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc=
125125
github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
126-
github.com/gptscript-ai/chat-completion-client v0.0.0-20240502162133-7dabc28eab59 h1:Nda0GDkrmIDiAFHfaXu0eIp6SsBs0/4Zyo87ff3Qcqo=
127-
github.com/gptscript-ai/chat-completion-client v0.0.0-20240502162133-7dabc28eab59/go.mod h1:7P/o6/IWa1KqsntVf68hSnLKuu3+xuqm6lYhch1w4jo=
126+
github.com/gptscript-ai/chat-completion-client v0.0.0-20240515050533-bdef9f2226a9 h1:s6nL/aokB1sJTqVXEjN0zFI5CJa66ubw9g68VTMzEw0=
127+
github.com/gptscript-ai/chat-completion-client v0.0.0-20240515050533-bdef9f2226a9/go.mod h1:7P/o6/IWa1KqsntVf68hSnLKuu3+xuqm6lYhch1w4jo=
128128
github.com/hashicorp/errwrap v1.0.0 h1:hLrqtEDnRye3+sgx6z4qVLNuviH3MR5aQ0ykNJa/UYA=
129129
github.com/hashicorp/errwrap v1.0.0/go.mod h1:YH+1FKiLXxHSkmPseP+kNlulaMuP3n2brvKWEqk/Jc4=
130130
github.com/hashicorp/go-multierror v1.1.1 h1:H5DkEtf6CXdFp0N0Em5UCwQpXMWke8IA0+lD48awMYo=

‎pkg/monitor/display.go‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,7 @@ type display struct {
6262
dumpState string
6363
callIDMap map[string]string
6464
callLock sync.Mutex
65+
usage types.Usage
6566
}
6667

6768
type livePrinter struct {
@@ -226,6 +227,10 @@ func (d *display) Event(event runner.Event) {
226227
userSpecifiedToolName: event.CallContext.ToolName,
227228
}
228229

230+
d.usage.PromptTokens += event.Usage.PromptTokens
231+
d.usage.CompletionTokens += event.Usage.CompletionTokens
232+
d.usage.TotalTokens += event.Usage.TotalTokens
233+
229234
switch event.Type {
230235
case runner.EventTypeCallStart:
231236
d.livePrinter.progressStart(currentCall)
@@ -283,6 +288,7 @@ func (d *display) Stop(output string, err error) {
283288
defer d.callLock.Unlock()
284289

285290
log.Fields("runID", d.dump.ID, "output", output, "err", err).Debugf("Run stopped")
291+
log.Fields("runID", d.dump.ID, "total", d.usage.TotalTokens, "prompt", d.usage.PromptTokens, "completion", d.usage.CompletionTokens).Infof("usage ")
286292
d.dump.Output = output
287293
d.dump.Err = err
288294
if d.dumpState != "" {

‎pkg/mvl/log.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,15 @@ func (f formatter) Format(entry *logrus.Entry) ([]byte, error) {
3838
if i, ok := entry.Data["response"]; ok && i != "" {
3939
msg += fmt.Sprintf(" [response=%s]", i)
4040
}
41+
if i, ok := entry.Data["total"]; ok && i != "" {
42+
msg += fmt.Sprintf(" [total=%v]", i)
43+
}
44+
if i, ok := entry.Data["prompt"]; ok && i != "" {
45+
msg += fmt.Sprintf(" [prompt=%v]", i)
46+
}
47+
if i, ok := entry.Data["completion"]; ok && i != "" {
48+
msg += fmt.Sprintf(" [completion=%v]", i)
49+
}
4150
return []byte(fmt.Sprintf("%s %s\n",
4251
entry.Time.Format(time.TimeOnly),
4352
msg)), nil

‎pkg/openai/client.go‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -308,6 +308,9 @@ func (c *Client) Call(ctx context.Context, messageRequest types.CompletionReques
308308
Model: messageRequest.Model,
309309
Messages: msgs,
310310
MaxTokens: messageRequest.MaxTokens,
311+
StreamOptions: &openai.StreamOptions{
312+
IncludeUsage: true,
313+
},
311314
}
312315

313316
if messageRequest.Temperature == nil {
@@ -372,17 +375,27 @@ func (c *Client) Call(ctx context.Context, messageRequest types.CompletionReques
372375
}
373376
}
374377

378+
var usage types.Usage
379+
if !cacheResponse {
380+
usage = result.Usage
381+
}
382+
375383
status <- types.CompletionStatus{
376384
CompletionID: id,
377385
Chunks: response,
378386
Response: result,
387+
Usage: usage,
379388
Cached: cacheResponse,
380389
}
381390

382391
return &result, nil
383392
}
384393

385394
func appendMessage(msg types.CompletionMessage, response openai.ChatCompletionStreamResponse) types.CompletionMessage {
395+
msg.Usage.CompletionTokens = types.FirstSet(msg.Usage.CompletionTokens, response.Usage.CompletionTokens)
396+
msg.Usage.PromptTokens = types.FirstSet(msg.Usage.PromptTokens, response.Usage.PromptTokens)
397+
msg.Usage.TotalTokens = types.FirstSet(msg.Usage.TotalTokens, response.Usage.TotalTokens)
398+
386399
if len(response.Choices) == 0 {
387400
return msg
388401
}
@@ -470,6 +483,7 @@ func (c *Client) call(ctx context.Context, request openai.ChatCompletionRequest,
470483
Object: resp.Object,
471484
Created: resp.Created,
472485
Model: resp.Model,
486+
Usage: resp.Usage,
473487
Choices: []openai.ChatCompletionStreamChoice{
474488
{
475489
Index: resp.Choices[0].Index,

‎pkg/runner/runner.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -196,6 +196,7 @@ type Event struct {
196196
ChatCompletionID string `json:"chatCompletionId,omitempty"`
197197
ChatRequest any `json:"chatRequest,omitempty"`
198198
ChatResponse any `json:"chatResponse,omitempty"`
199+
Usage types.Usage `json:"usage,omitempty"`
199200
ChatResponseCached bool `json:"chatResponseCached,omitempty"`
200201
Content string `json:"content,omitempty"`
201202
}
@@ -622,6 +623,7 @@ func streamProgress(callCtx *engine.Context, monitor Monitor) (chan<- types.Comp
622623
ChatCompletionID: status.CompletionID,
623624
ChatRequest: status.Request,
624625
ChatResponse: status.Response,
626+
Usage: status.Usage,
625627
ChatResponseCached: status.Cached,
626628
})
627629
}

‎pkg/tests/runner_test.go‎

Lines changed: 40 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -247,15 +247,17 @@ func TestSubChat(t *testing.T) {
247247
{
248248
"text": "Call chatbot"
249249
}
250-
]
250+
],
251+
"usage": {}
251252
},
252253
{
253254
"role": "user",
254255
"content": [
255256
{
256257
"text": "Hello"
257258
}
258-
]
259+
],
260+
"usage": {}
259261
},
260262
{
261263
"role": "assistant",
@@ -269,7 +271,8 @@ func TestSubChat(t *testing.T) {
269271
}
270272
}
271273
}
272-
]
274+
],
275+
"usage": {}
273276
}
274277
],
275278
"MaxTokens": 0,
@@ -312,15 +315,17 @@ func TestSubChat(t *testing.T) {
312315
{
313316
"text": "This is a chatbot"
314317
}
315-
]
318+
],
319+
"usage": {}
316320
},
317321
{
318322
"role": "assistant",
319323
"content": [
320324
{
321325
"text": "Assistant 1"
322326
}
323-
]
327+
],
328+
"usage": {}
324329
}
325330
],
326331
"MaxTokens": 0,
@@ -370,15 +375,17 @@ func TestSubChat(t *testing.T) {
370375
{
371376
"text": "Call chatbot"
372377
}
373-
]
378+
],
379+
"usage": {}
374380
},
375381
{
376382
"role": "user",
377383
"content": [
378384
{
379385
"text": "Hello"
380386
}
381-
]
387+
],
388+
"usage": {}
382389
},
383390
{
384391
"role": "assistant",
@@ -392,7 +399,8 @@ func TestSubChat(t *testing.T) {
392399
}
393400
}
394401
}
395-
]
402+
],
403+
"usage": {}
396404
}
397405
],
398406
"MaxTokens": 0,
@@ -435,31 +443,35 @@ func TestSubChat(t *testing.T) {
435443
{
436444
"text": "This is a chatbot"
437445
}
438-
]
446+
],
447+
"usage": {}
439448
},
440449
{
441450
"role": "assistant",
442451
"content": [
443452
{
444453
"text": "Assistant 1"
445454
}
446-
]
455+
],
456+
"usage": {}
447457
},
448458
{
449459
"role": "user",
450460
"content": [
451461
{
452462
"text": "User 1"
453463
}
454-
]
464+
],
465+
"usage": {}
455466
},
456467
{
457468
"role": "assistant",
458469
"content": [
459470
{
460471
"text": "Assistant 2"
461472
}
462-
]
473+
],
474+
"usage": {}
463475
}
464476
],
465477
"MaxTokens": 0,
@@ -513,23 +525,26 @@ func TestChat(t *testing.T) {
513525
{
514526
"text": "This is a chatbot"
515527
}
516-
]
528+
],
529+
"usage": {}
517530
},
518531
{
519532
"role": "user",
520533
"content": [
521534
{
522535
"text": "Hello"
523536
}
524-
]
537+
],
538+
"usage": {}
525539
},
526540
{
527541
"role": "assistant",
528542
"content": [
529543
{
530544
"text": "Assistant 1"
531545
}
532-
]
546+
],
547+
"usage": {}
533548
}
534549
],
535550
"MaxTokens": 0,
@@ -567,39 +582,44 @@ func TestChat(t *testing.T) {
567582
{
568583
"text": "This is a chatbot"
569584
}
570-
]
585+
],
586+
"usage": {}
571587
},
572588
{
573589
"role": "user",
574590
"content": [
575591
{
576592
"text": "Hello"
577593
}
578-
]
594+
],
595+
"usage": {}
579596
},
580597
{
581598
"role": "assistant",
582599
"content": [
583600
{
584601
"text": "Assistant 1"
585602
}
586-
]
603+
],
604+
"usage": {}
587605
},
588606
{
589607
"role": "user",
590608
"content": [
591609
{
592610
"text": "User 1"
593611
}
594-
]
612+
],
613+
"usage": {}
595614
},
596615
{
597616
"role": "assistant",
598617
"content": [
599618
{
600619
"text": "Assistant 2"
601620
}
602-
]
621+
],
622+
"usage": {}
603623
}
604624
],
605625
"MaxTokens": 0,

‎pkg/tests/testdata/TestCase/call1.golden‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,8 @@
2626
{
2727
"text": "Ask Bob how he is doing and let me know exactly what he said."
2828
}
29-
]
29+
],
30+
"usage": {}
3031
}
3132
],
3233
"MaxTokens": 0,

‎pkg/tests/testdata/TestCase2/call1.golden‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,8 @@
2626
{
2727
"text": "Ask Bob how he is doing and let me know exactly what he said."
2828
}
29-
]
29+
],
30+
"usage": {}
3031
}
3132
],
3233
"MaxTokens": 0,

‎pkg/tests/testdata/TestChat/call1.golden‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,15 +9,17 @@
99
{
1010
"text": "This is a chatbot"
1111
}
12-
]
12+
],
13+
"usage": {}
1314
},
1415
{
1516
"role": "user",
1617
"content": [
1718
{
1819
"text": "Hello"
1920
}
20-
]
21+
],
22+
"usage": {}
2123
}
2224
],
2325
"MaxTokens": 0,

0 commit comments

Comments
 (0)