From 71261aa4d6a31bbfce6333514f0136d1d6fd007d Mon Sep 17 00:00:00 2001 From: Darren Shepherd Date: Tue, 14 May 2024 22:13:04 -0700 Subject: [PATCH] chore: collect and log token usage --- go.mod | 2 +- go.sum | 4 +- pkg/monitor/display.go | 6 ++ pkg/mvl/log.go | 9 +++ pkg/openai/client.go | 14 +++++ pkg/runner/runner.go | 2 + pkg/tests/runner_test.go | 60 ++++++++++++------- pkg/tests/testdata/TestCase/call1.golden | 3 +- pkg/tests/testdata/TestCase2/call1.golden | 3 +- pkg/tests/testdata/TestChat/call1.golden | 6 +- pkg/tests/testdata/TestChat/call2.golden | 12 ++-- .../testdata/TestChatRunNoError/call1.golden | 3 +- pkg/tests/testdata/TestContext/call1.golden | 3 +- .../testdata/TestContextArg/call1.golden | 6 +- .../testdata/TestContextSubChat/call1.golden | 3 +- .../testdata/TestContextSubChat/call10.golden | 12 ++-- .../testdata/TestContextSubChat/call2.golden | 6 +- .../testdata/TestContextSubChat/call3.golden | 12 ++-- .../testdata/TestContextSubChat/call4.golden | 9 ++- .../testdata/TestContextSubChat/call5.golden | 6 +- .../testdata/TestContextSubChat/call6.golden | 3 +- .../testdata/TestContextSubChat/call7.golden | 6 +- .../testdata/TestContextSubChat/call8.golden | 12 ++-- .../testdata/TestContextSubChat/call9.golden | 9 ++- .../testdata/TestContextSubChat/step1.golden | 15 +++-- .../testdata/TestContextSubChat/step2.golden | 9 ++- .../testdata/TestContextSubChat/step3.golden | 24 +++++--- .../testdata/TestContextSubChat/step4.golden | 15 +++-- pkg/tests/testdata/TestCwd/call1.golden | 3 +- pkg/tests/testdata/TestCwd/call2.golden | 9 ++- pkg/tests/testdata/TestCwd/call3.golden | 15 +++-- .../testdata/TestDualSubChat/call1.golden | 6 +- .../testdata/TestDualSubChat/call2.golden | 6 +- .../testdata/TestDualSubChat/call3.golden | 6 +- .../testdata/TestDualSubChat/call4.golden | 12 ++-- .../testdata/TestDualSubChat/call5.golden | 12 ++-- .../testdata/TestDualSubChat/call6.golden | 18 ++++-- .../testdata/TestDualSubChat/call7.golden | 15 +++-- .../testdata/TestDualSubChat/step1.golden | 27 ++++++--- .../testdata/TestDualSubChat/step2.golden | 18 ++++-- .../testdata/TestDualSubChat/step3.golden | 24 +++++--- pkg/tests/testdata/TestExport/call1.golden | 3 +- pkg/tests/testdata/TestExport/call2.golden | 3 +- pkg/tests/testdata/TestExport/call3.golden | 9 ++- .../testdata/TestExportContext/call1.golden | 3 +- pkg/tests/testdata/TestSubChat/call1.golden | 6 +- pkg/tests/testdata/TestSubChat/call2.golden | 3 +- pkg/tests/testdata/TestSubChat/call3.golden | 9 ++- pkg/types/completion.go | 8 +++ 49 files changed, 338 insertions(+), 151 deletions(-) diff --git a/go.mod b/go.mod index 69dff731..1c2e595e 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index e3d56ef4..89161f17 100644 --- a/go.sum +++ b/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= diff --git a/pkg/monitor/display.go b/pkg/monitor/display.go index 51ac6470..c84e9cdb 100644 --- a/pkg/monitor/display.go +++ b/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 != "" { diff --git a/pkg/mvl/log.go b/pkg/mvl/log.go index 422d0644..a558ab34 100644 --- a/pkg/mvl/log.go +++ b/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 diff --git a/pkg/openai/client.go b/pkg/openai/client.go index 3f154d73..a04ce2d4 100644 --- a/pkg/openai/client.go +++ b/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,10 +375,16 @@ 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, } @@ -383,6 +392,10 @@ func (c *Client) Call(ctx context.Context, messageRequest types.CompletionReques } 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, diff --git a/pkg/runner/runner.go b/pkg/runner/runner.go index 67c3d652..b93511e1 100644 --- a/pkg/runner/runner.go +++ b/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, }) } diff --git a/pkg/tests/runner_test.go b/pkg/tests/runner_test.go index 700d965a..9941aad0 100644 --- a/pkg/tests/runner_test.go +++ b/pkg/tests/runner_test.go @@ -247,7 +247,8 @@ func TestSubChat(t *testing.T) { { "text": "Call chatbot" } - ] + ], + "usage": {} }, { "role": "user", @@ -255,7 +256,8 @@ func TestSubChat(t *testing.T) { { "text": "Hello" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -269,7 +271,8 @@ func TestSubChat(t *testing.T) { } } } - ] + ], + "usage": {} } ], "MaxTokens": 0, @@ -312,7 +315,8 @@ func TestSubChat(t *testing.T) { { "text": "This is a chatbot" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -320,7 +324,8 @@ func TestSubChat(t *testing.T) { { "text": "Assistant 1" } - ] + ], + "usage": {} } ], "MaxTokens": 0, @@ -370,7 +375,8 @@ func TestSubChat(t *testing.T) { { "text": "Call chatbot" } - ] + ], + "usage": {} }, { "role": "user", @@ -378,7 +384,8 @@ func TestSubChat(t *testing.T) { { "text": "Hello" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -392,7 +399,8 @@ func TestSubChat(t *testing.T) { } } } - ] + ], + "usage": {} } ], "MaxTokens": 0, @@ -435,7 +443,8 @@ func TestSubChat(t *testing.T) { { "text": "This is a chatbot" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -443,7 +452,8 @@ func TestSubChat(t *testing.T) { { "text": "Assistant 1" } - ] + ], + "usage": {} }, { "role": "user", @@ -451,7 +461,8 @@ func TestSubChat(t *testing.T) { { "text": "User 1" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -459,7 +470,8 @@ func TestSubChat(t *testing.T) { { "text": "Assistant 2" } - ] + ], + "usage": {} } ], "MaxTokens": 0, @@ -513,7 +525,8 @@ func TestChat(t *testing.T) { { "text": "This is a chatbot" } - ] + ], + "usage": {} }, { "role": "user", @@ -521,7 +534,8 @@ func TestChat(t *testing.T) { { "text": "Hello" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -529,7 +543,8 @@ func TestChat(t *testing.T) { { "text": "Assistant 1" } - ] + ], + "usage": {} } ], "MaxTokens": 0, @@ -567,7 +582,8 @@ func TestChat(t *testing.T) { { "text": "This is a chatbot" } - ] + ], + "usage": {} }, { "role": "user", @@ -575,7 +591,8 @@ func TestChat(t *testing.T) { { "text": "Hello" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -583,7 +600,8 @@ func TestChat(t *testing.T) { { "text": "Assistant 1" } - ] + ], + "usage": {} }, { "role": "user", @@ -591,7 +609,8 @@ func TestChat(t *testing.T) { { "text": "User 1" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -599,7 +618,8 @@ func TestChat(t *testing.T) { { "text": "Assistant 2" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestCase/call1.golden b/pkg/tests/testdata/TestCase/call1.golden index ab170df9..0ed42d15 100644 --- a/pkg/tests/testdata/TestCase/call1.golden +++ b/pkg/tests/testdata/TestCase/call1.golden @@ -26,7 +26,8 @@ { "text": "Ask Bob how he is doing and let me know exactly what he said." } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestCase2/call1.golden b/pkg/tests/testdata/TestCase2/call1.golden index 07d94234..d4998309 100644 --- a/pkg/tests/testdata/TestCase2/call1.golden +++ b/pkg/tests/testdata/TestCase2/call1.golden @@ -26,7 +26,8 @@ { "text": "Ask Bob how he is doing and let me know exactly what he said." } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestChat/call1.golden b/pkg/tests/testdata/TestChat/call1.golden index d9f9af0d..2f36d947 100644 --- a/pkg/tests/testdata/TestChat/call1.golden +++ b/pkg/tests/testdata/TestChat/call1.golden @@ -9,7 +9,8 @@ { "text": "This is a chatbot" } - ] + ], + "usage": {} }, { "role": "user", @@ -17,7 +18,8 @@ { "text": "Hello" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestChat/call2.golden b/pkg/tests/testdata/TestChat/call2.golden index 9a21703a..5a45cbfa 100644 --- a/pkg/tests/testdata/TestChat/call2.golden +++ b/pkg/tests/testdata/TestChat/call2.golden @@ -9,7 +9,8 @@ { "text": "This is a chatbot" } - ] + ], + "usage": {} }, { "role": "user", @@ -17,7 +18,8 @@ { "text": "Hello" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -25,7 +27,8 @@ { "text": "Assistant 1" } - ] + ], + "usage": {} }, { "role": "user", @@ -33,7 +36,8 @@ { "text": "User 1" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestChatRunNoError/call1.golden b/pkg/tests/testdata/TestChatRunNoError/call1.golden index ce624637..d0b56407 100644 --- a/pkg/tests/testdata/TestChatRunNoError/call1.golden +++ b/pkg/tests/testdata/TestChatRunNoError/call1.golden @@ -9,7 +9,8 @@ { "text": "This is a chatbot" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestContext/call1.golden b/pkg/tests/testdata/TestContext/call1.golden index 6833267b..356c871a 100644 --- a/pkg/tests/testdata/TestContext/call1.golden +++ b/pkg/tests/testdata/TestContext/call1.golden @@ -9,7 +9,8 @@ { "text": "this is from context\n\nThis is from tool" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestContextArg/call1.golden b/pkg/tests/testdata/TestContextArg/call1.golden index cc4213fb..e51539cc 100644 --- a/pkg/tests/testdata/TestContextArg/call1.golden +++ b/pkg/tests/testdata/TestContextArg/call1.golden @@ -9,7 +9,8 @@ { "text": "this is from context -- foo.db\n\nthis is from other context foo.db and then\n\nthis is from other context and then foo.db\n\nThis is from tool" } - ] + ], + "usage": {} }, { "role": "user", @@ -17,7 +18,8 @@ { "text": "{\n\"file\": \"foo.db\"\n}" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestContextSubChat/call1.golden b/pkg/tests/testdata/TestContextSubChat/call1.golden index fe35d630..5bd16621 100644 --- a/pkg/tests/testdata/TestContextSubChat/call1.golden +++ b/pkg/tests/testdata/TestContextSubChat/call1.golden @@ -17,7 +17,8 @@ { "text": "Call chatbot" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestContextSubChat/call10.golden b/pkg/tests/testdata/TestContextSubChat/call10.golden index aabe5587..236dded1 100644 --- a/pkg/tests/testdata/TestContextSubChat/call10.golden +++ b/pkg/tests/testdata/TestContextSubChat/call10.golden @@ -9,7 +9,8 @@ { "text": "Assistant Response 5 - from context tool resume\nHello" } - ] + ], + "usage": {} }, { "role": "user", @@ -17,7 +18,8 @@ { "text": "User 1" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -25,7 +27,8 @@ { "text": "Assistant Response 3 - from main chat tool" } - ] + ], + "usage": {} }, { "role": "user", @@ -33,7 +36,8 @@ { "text": "User 3" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestContextSubChat/call2.golden b/pkg/tests/testdata/TestContextSubChat/call2.golden index 2d29bfc8..16cd2012 100644 --- a/pkg/tests/testdata/TestContextSubChat/call2.golden +++ b/pkg/tests/testdata/TestContextSubChat/call2.golden @@ -26,7 +26,8 @@ { "text": "This is a chatbot" } - ] + ], + "usage": {} }, { "role": "user", @@ -34,7 +35,8 @@ { "text": "Input to chatbot1" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestContextSubChat/call3.golden b/pkg/tests/testdata/TestContextSubChat/call3.golden index cf2a4c29..e46dae91 100644 --- a/pkg/tests/testdata/TestContextSubChat/call3.golden +++ b/pkg/tests/testdata/TestContextSubChat/call3.golden @@ -26,7 +26,8 @@ { "text": "This is a chatbot" } - ] + ], + "usage": {} }, { "role": "user", @@ -34,7 +35,8 @@ { "text": "Input to chatbot1" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -42,7 +44,8 @@ { "text": "Assistant Response 1 - from chatbot1" } - ] + ], + "usage": {} }, { "role": "user", @@ -50,7 +53,8 @@ { "text": "User 1" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestContextSubChat/call4.golden b/pkg/tests/testdata/TestContextSubChat/call4.golden index 98d932af..78de5261 100644 --- a/pkg/tests/testdata/TestContextSubChat/call4.golden +++ b/pkg/tests/testdata/TestContextSubChat/call4.golden @@ -17,7 +17,8 @@ { "text": "Call chatbot" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -32,7 +33,8 @@ } } } - ] + ], + "usage": {} }, { "role": "tool", @@ -48,7 +50,8 @@ "name": "chatbot", "arguments": "Input to chatbot1" } - } + }, + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestContextSubChat/call5.golden b/pkg/tests/testdata/TestContextSubChat/call5.golden index 974ff85f..ae69a80a 100644 --- a/pkg/tests/testdata/TestContextSubChat/call5.golden +++ b/pkg/tests/testdata/TestContextSubChat/call5.golden @@ -9,7 +9,8 @@ { "text": "Assistant Response 2 - from context tool\nHello" } - ] + ], + "usage": {} }, { "role": "user", @@ -17,7 +18,8 @@ { "text": "User 1" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestContextSubChat/call6.golden b/pkg/tests/testdata/TestContextSubChat/call6.golden index fe35d630..5bd16621 100644 --- a/pkg/tests/testdata/TestContextSubChat/call6.golden +++ b/pkg/tests/testdata/TestContextSubChat/call6.golden @@ -17,7 +17,8 @@ { "text": "Call chatbot" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestContextSubChat/call7.golden b/pkg/tests/testdata/TestContextSubChat/call7.golden index d12c4297..047255b3 100644 --- a/pkg/tests/testdata/TestContextSubChat/call7.golden +++ b/pkg/tests/testdata/TestContextSubChat/call7.golden @@ -26,7 +26,8 @@ { "text": "This is a chatbot" } - ] + ], + "usage": {} }, { "role": "user", @@ -34,7 +35,8 @@ { "text": "Input to chatbot1 on resume" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestContextSubChat/call8.golden b/pkg/tests/testdata/TestContextSubChat/call8.golden index 8cbffd87..3e9ec556 100644 --- a/pkg/tests/testdata/TestContextSubChat/call8.golden +++ b/pkg/tests/testdata/TestContextSubChat/call8.golden @@ -26,7 +26,8 @@ { "text": "This is a chatbot" } - ] + ], + "usage": {} }, { "role": "user", @@ -34,7 +35,8 @@ { "text": "Input to chatbot1 on resume" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -42,7 +44,8 @@ { "text": "Assistant Response 4 - from chatbot1" } - ] + ], + "usage": {} }, { "role": "user", @@ -50,7 +53,8 @@ { "text": "User 4" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestContextSubChat/call9.golden b/pkg/tests/testdata/TestContextSubChat/call9.golden index f45e0e01..9572de06 100644 --- a/pkg/tests/testdata/TestContextSubChat/call9.golden +++ b/pkg/tests/testdata/TestContextSubChat/call9.golden @@ -17,7 +17,8 @@ { "text": "Call chatbot" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -32,7 +33,8 @@ } } } - ] + ], + "usage": {} }, { "role": "tool", @@ -48,7 +50,8 @@ "name": "chatbot", "arguments": "Input to chatbot1 on resume" } - } + }, + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestContextSubChat/step1.golden b/pkg/tests/testdata/TestContextSubChat/step1.golden index 53ef8e7b..2c345ed1 100644 --- a/pkg/tests/testdata/TestContextSubChat/step1.golden +++ b/pkg/tests/testdata/TestContextSubChat/step1.golden @@ -25,7 +25,8 @@ { "text": "Call chatbot" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -40,7 +41,8 @@ } } } - ] + ], + "usage": {} } ], "MaxTokens": 0, @@ -103,7 +105,8 @@ { "text": "This is a chatbot" } - ] + ], + "usage": {} }, { "role": "user", @@ -111,7 +114,8 @@ { "text": "Input to chatbot1" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -119,7 +123,8 @@ { "text": "Assistant Response 1 - from chatbot1" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestContextSubChat/step2.golden b/pkg/tests/testdata/TestContextSubChat/step2.golden index 5e56f4f9..7394a502 100644 --- a/pkg/tests/testdata/TestContextSubChat/step2.golden +++ b/pkg/tests/testdata/TestContextSubChat/step2.golden @@ -17,7 +17,8 @@ { "text": "Assistant Response 2 - from context tool\nHello" } - ] + ], + "usage": {} }, { "role": "user", @@ -25,7 +26,8 @@ { "text": "User 1" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -33,7 +35,8 @@ { "text": "Assistant Response 3 - from main chat tool" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestContextSubChat/step3.golden b/pkg/tests/testdata/TestContextSubChat/step3.golden index e8fbfcfb..3d2dcf90 100644 --- a/pkg/tests/testdata/TestContextSubChat/step3.golden +++ b/pkg/tests/testdata/TestContextSubChat/step3.golden @@ -17,7 +17,8 @@ { "text": "Assistant Response 2 - from context tool\nHello" } - ] + ], + "usage": {} }, { "role": "user", @@ -25,7 +26,8 @@ { "text": "User 1" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -33,7 +35,8 @@ { "text": "Assistant Response 3 - from main chat tool" } - ] + ], + "usage": {} } ], "MaxTokens": 0, @@ -69,7 +72,8 @@ { "text": "Call chatbot" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -84,7 +88,8 @@ } } } - ] + ], + "usage": {} } ], "MaxTokens": 0, @@ -147,7 +152,8 @@ { "text": "This is a chatbot" } - ] + ], + "usage": {} }, { "role": "user", @@ -155,7 +161,8 @@ { "text": "Input to chatbot1 on resume" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -163,7 +170,8 @@ { "text": "Assistant Response 4 - from chatbot1" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestContextSubChat/step4.golden b/pkg/tests/testdata/TestContextSubChat/step4.golden index 882eeecf..6d975e1c 100644 --- a/pkg/tests/testdata/TestContextSubChat/step4.golden +++ b/pkg/tests/testdata/TestContextSubChat/step4.golden @@ -17,7 +17,8 @@ { "text": "Assistant Response 5 - from context tool resume\nHello" } - ] + ], + "usage": {} }, { "role": "user", @@ -25,7 +26,8 @@ { "text": "User 1" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -33,7 +35,8 @@ { "text": "Assistant Response 3 - from main chat tool" } - ] + ], + "usage": {} }, { "role": "user", @@ -41,7 +44,8 @@ { "text": "User 3" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -49,7 +53,8 @@ { "text": "Assistant Response 6 - from main chat tool resume" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestCwd/call1.golden b/pkg/tests/testdata/TestCwd/call1.golden index 4cc327c7..f366cc42 100644 --- a/pkg/tests/testdata/TestCwd/call1.golden +++ b/pkg/tests/testdata/TestCwd/call1.golden @@ -24,7 +24,8 @@ { "text": "noop" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestCwd/call2.golden b/pkg/tests/testdata/TestCwd/call2.golden index afd81594..f61d473a 100644 --- a/pkg/tests/testdata/TestCwd/call2.golden +++ b/pkg/tests/testdata/TestCwd/call2.golden @@ -24,7 +24,8 @@ { "text": "noop" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -38,7 +39,8 @@ } } } - ] + ], + "usage": {} }, { "role": "tool", @@ -53,7 +55,8 @@ "function": { "name": "test" } - } + }, + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestCwd/call3.golden b/pkg/tests/testdata/TestCwd/call3.golden index d1a485ab..3e310b7e 100644 --- a/pkg/tests/testdata/TestCwd/call3.golden +++ b/pkg/tests/testdata/TestCwd/call3.golden @@ -24,7 +24,8 @@ { "text": "noop" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -38,7 +39,8 @@ } } } - ] + ], + "usage": {} }, { "role": "tool", @@ -53,7 +55,8 @@ "function": { "name": "test" } - } + }, + "usage": {} }, { "role": "assistant", @@ -67,7 +70,8 @@ } } } - ] + ], + "usage": {} }, { "role": "tool", @@ -82,7 +86,8 @@ "function": { "name": "local" } - } + }, + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestDualSubChat/call1.golden b/pkg/tests/testdata/TestDualSubChat/call1.golden index 529f6f68..852fe054 100644 --- a/pkg/tests/testdata/TestDualSubChat/call1.golden +++ b/pkg/tests/testdata/TestDualSubChat/call1.golden @@ -24,7 +24,8 @@ { "text": "Call chatbots" } - ] + ], + "usage": {} }, { "role": "user", @@ -32,7 +33,8 @@ { "text": "User 1" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestDualSubChat/call2.golden b/pkg/tests/testdata/TestDualSubChat/call2.golden index 2d29bfc8..16cd2012 100644 --- a/pkg/tests/testdata/TestDualSubChat/call2.golden +++ b/pkg/tests/testdata/TestDualSubChat/call2.golden @@ -26,7 +26,8 @@ { "text": "This is a chatbot" } - ] + ], + "usage": {} }, { "role": "user", @@ -34,7 +35,8 @@ { "text": "Input to chatbot1" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestDualSubChat/call3.golden b/pkg/tests/testdata/TestDualSubChat/call3.golden index 415c8332..7ed7ae76 100644 --- a/pkg/tests/testdata/TestDualSubChat/call3.golden +++ b/pkg/tests/testdata/TestDualSubChat/call3.golden @@ -26,7 +26,8 @@ { "text": "This is a chatbot2" } - ] + ], + "usage": {} }, { "role": "user", @@ -34,7 +35,8 @@ { "text": "Input to chatbot2" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestDualSubChat/call4.golden b/pkg/tests/testdata/TestDualSubChat/call4.golden index f1fa4ecf..ab2aace9 100644 --- a/pkg/tests/testdata/TestDualSubChat/call4.golden +++ b/pkg/tests/testdata/TestDualSubChat/call4.golden @@ -26,7 +26,8 @@ { "text": "This is a chatbot" } - ] + ], + "usage": {} }, { "role": "user", @@ -34,7 +35,8 @@ { "text": "Input to chatbot1" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -42,7 +44,8 @@ { "text": "Assistant Response 1 - from chatbot1" } - ] + ], + "usage": {} }, { "role": "user", @@ -50,7 +53,8 @@ { "text": "User 2" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestDualSubChat/call5.golden b/pkg/tests/testdata/TestDualSubChat/call5.golden index 165a064b..6b9cb0d8 100644 --- a/pkg/tests/testdata/TestDualSubChat/call5.golden +++ b/pkg/tests/testdata/TestDualSubChat/call5.golden @@ -26,7 +26,8 @@ { "text": "This is a chatbot2" } - ] + ], + "usage": {} }, { "role": "user", @@ -34,7 +35,8 @@ { "text": "Input to chatbot2" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -42,7 +44,8 @@ { "text": "Assistent Response 2 - from chatbot2" } - ] + ], + "usage": {} }, { "role": "user", @@ -50,7 +53,8 @@ { "text": "User 3" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestDualSubChat/call6.golden b/pkg/tests/testdata/TestDualSubChat/call6.golden index 69bf9804..b9f6e3d3 100644 --- a/pkg/tests/testdata/TestDualSubChat/call6.golden +++ b/pkg/tests/testdata/TestDualSubChat/call6.golden @@ -26,7 +26,8 @@ { "text": "This is a chatbot2" } - ] + ], + "usage": {} }, { "role": "user", @@ -34,7 +35,8 @@ { "text": "Input to chatbot2" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -42,7 +44,8 @@ { "text": "Assistent Response 2 - from chatbot2" } - ] + ], + "usage": {} }, { "role": "user", @@ -50,7 +53,8 @@ { "text": "User 3" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -58,7 +62,8 @@ { "text": "Assistant 3" } - ] + ], + "usage": {} }, { "role": "user", @@ -66,7 +71,8 @@ { "text": "User 4" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestDualSubChat/call7.golden b/pkg/tests/testdata/TestDualSubChat/call7.golden index 2f98f090..31696c01 100644 --- a/pkg/tests/testdata/TestDualSubChat/call7.golden +++ b/pkg/tests/testdata/TestDualSubChat/call7.golden @@ -24,7 +24,8 @@ { "text": "Call chatbots" } - ] + ], + "usage": {} }, { "role": "user", @@ -32,7 +33,8 @@ { "text": "User 1" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -57,7 +59,8 @@ } } } - ] + ], + "usage": {} }, { "role": "tool", @@ -73,7 +76,8 @@ "name": "chatbot", "arguments": "Input to chatbot1" } - } + }, + "usage": {} }, { "role": "tool", @@ -89,7 +93,8 @@ "name": "chatbot2", "arguments": "Input to chatbot2" } - } + }, + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestDualSubChat/step1.golden b/pkg/tests/testdata/TestDualSubChat/step1.golden index c02e67a9..ab54e4b8 100644 --- a/pkg/tests/testdata/TestDualSubChat/step1.golden +++ b/pkg/tests/testdata/TestDualSubChat/step1.golden @@ -32,7 +32,8 @@ { "text": "Call chatbots" } - ] + ], + "usage": {} }, { "role": "user", @@ -40,7 +41,8 @@ { "text": "User 1" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -65,7 +67,8 @@ } } } - ] + ], + "usage": {} } ], "MaxTokens": 0, @@ -140,7 +143,8 @@ { "text": "This is a chatbot" } - ] + ], + "usage": {} }, { "role": "user", @@ -148,7 +152,8 @@ { "text": "Input to chatbot1" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -156,7 +161,8 @@ { "text": "Assistant Response 1 - from chatbot1" } - ] + ], + "usage": {} } ], "MaxTokens": 0, @@ -206,7 +212,8 @@ { "text": "This is a chatbot2" } - ] + ], + "usage": {} }, { "role": "user", @@ -214,7 +221,8 @@ { "text": "Input to chatbot2" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -222,7 +230,8 @@ { "text": "Assistent Response 2 - from chatbot2" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestDualSubChat/step2.golden b/pkg/tests/testdata/TestDualSubChat/step2.golden index 0f090762..3575765b 100644 --- a/pkg/tests/testdata/TestDualSubChat/step2.golden +++ b/pkg/tests/testdata/TestDualSubChat/step2.golden @@ -32,7 +32,8 @@ { "text": "Call chatbots" } - ] + ], + "usage": {} }, { "role": "user", @@ -40,7 +41,8 @@ { "text": "User 1" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -65,7 +67,8 @@ } } } - ] + ], + "usage": {} } ], "MaxTokens": 0, @@ -147,7 +150,8 @@ { "text": "This is a chatbot2" } - ] + ], + "usage": {} }, { "role": "user", @@ -155,7 +159,8 @@ { "text": "Input to chatbot2" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -163,7 +168,8 @@ { "text": "Assistent Response 2 - from chatbot2" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestDualSubChat/step3.golden b/pkg/tests/testdata/TestDualSubChat/step3.golden index a35becd7..c25c1883 100644 --- a/pkg/tests/testdata/TestDualSubChat/step3.golden +++ b/pkg/tests/testdata/TestDualSubChat/step3.golden @@ -32,7 +32,8 @@ { "text": "Call chatbots" } - ] + ], + "usage": {} }, { "role": "user", @@ -40,7 +41,8 @@ { "text": "User 1" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -65,7 +67,8 @@ } } } - ] + ], + "usage": {} } ], "MaxTokens": 0, @@ -147,7 +150,8 @@ { "text": "This is a chatbot2" } - ] + ], + "usage": {} }, { "role": "user", @@ -155,7 +159,8 @@ { "text": "Input to chatbot2" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -163,7 +168,8 @@ { "text": "Assistent Response 2 - from chatbot2" } - ] + ], + "usage": {} }, { "role": "user", @@ -171,7 +177,8 @@ { "text": "User 3" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -179,7 +186,8 @@ { "text": "Assistant 3" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestExport/call1.golden b/pkg/tests/testdata/TestExport/call1.golden index d46fd334..852208e0 100644 --- a/pkg/tests/testdata/TestExport/call1.golden +++ b/pkg/tests/testdata/TestExport/call1.golden @@ -64,7 +64,8 @@ { "text": "the default" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestExport/call2.golden b/pkg/tests/testdata/TestExport/call2.golden index ffb73ff4..795c4dcd 100644 --- a/pkg/tests/testdata/TestExport/call2.golden +++ b/pkg/tests/testdata/TestExport/call2.golden @@ -9,7 +9,8 @@ { "text": "the transient" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestExport/call3.golden b/pkg/tests/testdata/TestExport/call3.golden index 632e8f44..16232d88 100644 --- a/pkg/tests/testdata/TestExport/call3.golden +++ b/pkg/tests/testdata/TestExport/call3.golden @@ -64,7 +64,8 @@ { "text": "the default" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -78,7 +79,8 @@ } } } - ] + ], + "usage": {} }, { "role": "tool", @@ -93,7 +95,8 @@ "function": { "name": "transient" } - } + }, + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestExportContext/call1.golden b/pkg/tests/testdata/TestExportContext/call1.golden index a93fbfcd..20add018 100644 --- a/pkg/tests/testdata/TestExportContext/call1.golden +++ b/pkg/tests/testdata/TestExportContext/call1.golden @@ -44,7 +44,8 @@ { "text": "this is from external context\n\nthis is from context\n\nThis is from tool" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestSubChat/call1.golden b/pkg/tests/testdata/TestSubChat/call1.golden index 6b05309e..0fee8b2a 100644 --- a/pkg/tests/testdata/TestSubChat/call1.golden +++ b/pkg/tests/testdata/TestSubChat/call1.golden @@ -17,7 +17,8 @@ { "text": "Call chatbot" } - ] + ], + "usage": {} }, { "role": "user", @@ -25,7 +26,8 @@ { "text": "Hello" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestSubChat/call2.golden b/pkg/tests/testdata/TestSubChat/call2.golden index ce624637..d0b56407 100644 --- a/pkg/tests/testdata/TestSubChat/call2.golden +++ b/pkg/tests/testdata/TestSubChat/call2.golden @@ -9,7 +9,8 @@ { "text": "This is a chatbot" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/tests/testdata/TestSubChat/call3.golden b/pkg/tests/testdata/TestSubChat/call3.golden index ac050849..1b5dfd1a 100644 --- a/pkg/tests/testdata/TestSubChat/call3.golden +++ b/pkg/tests/testdata/TestSubChat/call3.golden @@ -9,7 +9,8 @@ { "text": "This is a chatbot" } - ] + ], + "usage": {} }, { "role": "assistant", @@ -17,7 +18,8 @@ { "text": "Assistant 1" } - ] + ], + "usage": {} }, { "role": "user", @@ -25,7 +27,8 @@ { "text": "User 1" } - ] + ], + "usage": {} } ], "MaxTokens": 0, diff --git a/pkg/types/completion.go b/pkg/types/completion.go index 48c18baf..35251d98 100644 --- a/pkg/types/completion.go +++ b/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