Skip to content

Commit 73494ba

Browse files
committed
fix: capture Usage, ChatResponseCached, and ToolResults
Additionally add tests to ensure these are captured properly. Signed-off-by: Donnie Adams <[email protected]>
1 parent 64eaa0a commit 73494ba

File tree

2 files changed

+80
-11
lines changed

2 files changed

+80
-11
lines changed

frame.go

+10-8
Original file line numberDiff line numberDiff line change
@@ -52,14 +52,16 @@ type RunFrame struct {
5252
type CallFrame struct {
5353
CallContext `json:",inline"`
5454

55-
Type EventType `json:"type"`
56-
Start time.Time `json:"start"`
57-
End time.Time `json:"end"`
58-
Input string `json:"input"`
59-
Output []Output `json:"output"`
60-
Usage Usage `json:"usage"`
61-
LLMRequest any `json:"llmRequest"`
62-
LLMResponse any `json:"llmResponse"`
55+
Type EventType `json:"type"`
56+
Start time.Time `json:"start"`
57+
End time.Time `json:"end"`
58+
Input string `json:"input"`
59+
Output []Output `json:"output"`
60+
Usage Usage `json:"usage"`
61+
ChatResponseCached bool `json:"chatResponseCached"`
62+
ToolResults int `json:"toolResults"`
63+
LLMRequest any `json:"llmRequest"`
64+
LLMResponse any `json:"llmResponse"`
6365
}
6466

6567
type Usage struct {

gptscript_test.go

+70-3
Original file line numberDiff line numberDiff line change
@@ -161,9 +161,7 @@ func TestAbortRun(t *testing.T) {
161161
func TestSimpleEvaluate(t *testing.T) {
162162
tool := ToolDef{Instructions: "What is the capital of the united states?"}
163163

164-
run, err := g.Evaluate(context.Background(), Options{
165-
GlobalOptions: GlobalOptions{},
166-
}, tool)
164+
run, err := g.Evaluate(context.Background(), Options{DisableCache: true}, tool)
167165
if err != nil {
168166
t.Errorf("Error executing tool: %v", err)
169167
}
@@ -190,6 +188,17 @@ func TestSimpleEvaluate(t *testing.T) {
190188
if run.Program() == nil {
191189
t.Error("Run program not set")
192190
}
191+
192+
var promptTokens, completionTokens, totalTokens int
193+
for _, c := range run.calls {
194+
promptTokens += c.Usage.PromptTokens
195+
completionTokens += c.Usage.CompletionTokens
196+
totalTokens += c.Usage.TotalTokens
197+
}
198+
199+
if promptTokens == 0 || completionTokens == 0 || totalTokens == 0 {
200+
t.Errorf("Usage not set: %d, %d, %d", promptTokens, completionTokens, totalTokens)
201+
}
193202
}
194203

195204
func TestEvaluateWithContext(t *testing.T) {
@@ -285,6 +294,16 @@ func TestEvaluateWithToolList(t *testing.T) {
285294
if !strings.Contains(out, "hello there") {
286295
t.Errorf("Unexpected output: %s", out)
287296
}
297+
298+
// In this case, we expect the total number of tool results to be 1
299+
var toolResults int
300+
for _, c := range run.calls {
301+
toolResults += c.ToolResults
302+
}
303+
304+
if toolResults != 1 {
305+
t.Errorf("Unexpected number of tool results: %d", toolResults)
306+
}
288307
}
289308

290309
func TestEvaluateWithToolListAndSubTool(t *testing.T) {
@@ -361,6 +380,54 @@ func TestStreamEvaluate(t *testing.T) {
361380
}
362381
}
363382

383+
func TestSimpleRun(t *testing.T) {
384+
wd, err := os.Getwd()
385+
if err != nil {
386+
t.Fatalf("Error getting working directory: %v", err)
387+
}
388+
389+
run, err := g.Run(context.Background(), wd+"/test/catcher.gpt", Options{})
390+
if err != nil {
391+
t.Fatalf("Error executing file: %v", err)
392+
}
393+
394+
out, err := run.Text()
395+
if err != nil {
396+
t.Errorf("Error reading output: %v", err)
397+
}
398+
399+
if !strings.Contains(out, "Salinger") {
400+
t.Errorf("Unexpected output: %s", out)
401+
}
402+
403+
if len(run.ErrorOutput()) != 0 {
404+
t.Error("Should have no stderr output")
405+
}
406+
407+
// Run it a second time, ensuring the same output and that a cached response is used
408+
run, err = g.Run(context.Background(), wd+"/test/catcher.gpt", Options{})
409+
if err != nil {
410+
t.Fatalf("Error executing file: %v", err)
411+
}
412+
413+
secondOut, err := run.Text()
414+
if err != nil {
415+
t.Errorf("Error reading output: %v", err)
416+
}
417+
418+
if secondOut != out {
419+
t.Errorf("Unexpected output on second run: %s != %s", out, secondOut)
420+
}
421+
422+
// In this case, we expect a single call and that the response is cached
423+
for _, c := range run.calls {
424+
if !c.ChatResponseCached {
425+
t.Error("Chat response should be cached")
426+
}
427+
break
428+
}
429+
}
430+
364431
func TestStreamRun(t *testing.T) {
365432
wd, err := os.Getwd()
366433
if err != nil {

0 commit comments

Comments
 (0)