Skip to content

Commit cee4564

Browse files
committed
replace SessionDigest.Tokens* with SessionDigest.UsageData
1 parent e8eb58a commit cee4564

16 files changed

Lines changed: 110 additions & 117 deletions

internal/baseline/baseline.go

Lines changed: 12 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -38,17 +38,21 @@ func ComputeImprovement(baseline, withSkill *models.RunResult) (float64, Improve
3838
breakdown.QualityDelta = skillScore - baselineScore
3939

4040
// Token reduction: negative means fewer tokens (good)
41-
baseTok := float64(baseline.SessionDigest.TokensTotal)
42-
skillTok := float64(withSkill.SessionDigest.TokensTotal)
43-
if baseTok > 0 {
44-
breakdown.TokenReduction = (skillTok - baseTok) / baseTok
41+
if baseline.SessionDigest.Usage != nil && withSkill.SessionDigest.Usage != nil {
42+
baseTok := float64(baseline.SessionDigest.Usage.InputTokens + baseline.SessionDigest.Usage.OutputTokens)
43+
skillTok := float64(withSkill.SessionDigest.Usage.InputTokens + withSkill.SessionDigest.Usage.OutputTokens)
44+
if baseTok > 0 {
45+
breakdown.TokenReduction = (skillTok - baseTok) / baseTok
46+
}
4547
}
4648

4749
// Turn reduction: negative means fewer turns (good)
48-
baseTurns := float64(baseline.SessionDigest.TotalTurns)
49-
skillTurns := float64(withSkill.SessionDigest.TotalTurns)
50-
if baseTurns > 0 {
51-
breakdown.TurnReduction = (skillTurns - baseTurns) / baseTurns
50+
if baseline.SessionDigest.Usage != nil && withSkill.SessionDigest.Usage != nil {
51+
baseTurns := float64(baseline.SessionDigest.Usage.Turns)
52+
skillTurns := float64(withSkill.SessionDigest.Usage.Turns)
53+
if baseTurns > 0 {
54+
breakdown.TurnReduction = (skillTurns - baseTurns) / baseTurns
55+
}
5256
}
5357

5458
// Time reduction: negative means faster (good)

internal/baseline/baseline_test.go

Lines changed: 11 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -15,8 +15,7 @@ func TestComputeImprovement_SkillBetter(t *testing.T) {
1515
"g1": {Score: 0.3, Weight: 1.0},
1616
},
1717
SessionDigest: models.SessionDigest{
18-
TotalTurns: 10,
19-
TokensTotal: 5000,
18+
Usage: &models.UsageStats{Turns: 10, InputTokens: 5000},
2019
},
2120
}
2221
withSkill := &models.RunResult{
@@ -26,8 +25,7 @@ func TestComputeImprovement_SkillBetter(t *testing.T) {
2625
"g1": {Score: 0.9, Weight: 1.0},
2726
},
2827
SessionDigest: models.SessionDigest{
29-
TotalTurns: 6,
30-
TokensTotal: 3000,
28+
Usage: &models.UsageStats{Turns: 6, InputTokens: 3000},
3129
},
3230
}
3331

@@ -49,8 +47,7 @@ func TestComputeImprovement_SkillWorse(t *testing.T) {
4947
"g1": {Score: 0.9, Weight: 1.0},
5048
},
5149
SessionDigest: models.SessionDigest{
52-
TotalTurns: 5,
53-
TokensTotal: 2000,
50+
Usage: &models.UsageStats{Turns: 5, InputTokens: 2000},
5451
},
5552
}
5653
withSkill := &models.RunResult{
@@ -60,8 +57,7 @@ func TestComputeImprovement_SkillWorse(t *testing.T) {
6057
"g1": {Score: 0.2, Weight: 1.0},
6158
},
6259
SessionDigest: models.SessionDigest{
63-
TotalTurns: 15,
64-
TokensTotal: 8000,
60+
Usage: &models.UsageStats{Turns: 15, InputTokens: 8000},
6561
},
6662
}
6763

@@ -83,8 +79,7 @@ func TestComputeImprovement_Equal(t *testing.T) {
8379
"g1": {Score: 0.8, Weight: 1.0},
8480
},
8581
SessionDigest: models.SessionDigest{
86-
TotalTurns: 5,
87-
TokensTotal: 2000,
82+
Usage: &models.UsageStats{Turns: 5, InputTokens: 2000},
8883
},
8984
}
9085

@@ -106,8 +101,7 @@ func TestComputeImprovement_ZeroBaseline(t *testing.T) {
106101
"g1": {Score: 0.0, Weight: 1.0},
107102
},
108103
SessionDigest: models.SessionDigest{
109-
TotalTurns: 0,
110-
TokensTotal: 0,
104+
Usage: &models.UsageStats{Turns: 0, InputTokens: 0},
111105
},
112106
}
113107
withSkill := &models.RunResult{
@@ -117,8 +111,7 @@ func TestComputeImprovement_ZeroBaseline(t *testing.T) {
117111
"g1": {Score: 1.0, Weight: 1.0},
118112
},
119113
SessionDigest: models.SessionDigest{
120-
TotalTurns: 5,
121-
TokensTotal: 3000,
114+
Usage: &models.UsageStats{Turns: 5, InputTokens: 3000},
122115
},
123116
}
124117

@@ -142,8 +135,7 @@ func TestComputeImprovement_ClampedToRange(t *testing.T) {
142135
"g1": {Score: 0.0, Weight: 1.0},
143136
},
144137
SessionDigest: models.SessionDigest{
145-
TotalTurns: 1,
146-
TokensTotal: 100,
138+
Usage: &models.UsageStats{Turns: 1, InputTokens: 100},
147139
},
148140
}
149141
withSkill := &models.RunResult{
@@ -153,8 +145,7 @@ func TestComputeImprovement_ClampedToRange(t *testing.T) {
153145
"g1": {Score: 1.0, Weight: 1.0},
154146
},
155147
SessionDigest: models.SessionDigest{
156-
TotalTurns: 1,
157-
TokensTotal: 10,
148+
Usage: &models.UsageStats{Turns: 1, InputTokens: 10},
158149
},
159150
}
160151

@@ -189,7 +180,7 @@ func TestComputeFromOutcomes_WithRuns(t *testing.T) {
189180
"g1": {Score: 0.9, Weight: 1.0},
190181
},
191182
SessionDigest: models.SessionDigest{
192-
TotalTurns: 5, TokensTotal: 2000,
183+
Usage: &models.UsageStats{Turns: 5, InputTokens: 2000},
193184
},
194185
},
195186
},
@@ -204,7 +195,7 @@ func TestComputeFromOutcomes_WithRuns(t *testing.T) {
204195
"g1": {Score: 0.3, Weight: 1.0},
205196
},
206197
SessionDigest: models.SessionDigest{
207-
TotalTurns: 10, TokensTotal: 5000,
198+
Usage: &models.UsageStats{Turns: 10, InputTokens: 5000},
208199
},
209200
},
210201
},

internal/execution/usage.go

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,6 @@ func UpdateOutcomeUsage(outcome *models.EvaluationOutcome, engine AgentEngine) {
1818
}
1919
if usage := engine.SessionUsage(run.SessionDigest.SessionID); usage != nil {
2020
run.SessionDigest.Usage = usage
21-
run.SessionDigest.TokensIn = usage.InputTokens
22-
run.SessionDigest.TokensOut = usage.OutputTokens
23-
run.SessionDigest.TokensTotal = usage.InputTokens + usage.OutputTokens
2421
}
2522
}
2623
}

internal/execution/usage_test.go

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -25,9 +25,6 @@ func TestUpdateOutcomeUsage_ReplacesAndReaggregates(t *testing.T) {
2525
InputTokens: 100,
2626
OutputTokens: 50,
2727
},
28-
TokensIn: 100,
29-
TokensOut: 50,
30-
TokensTotal: 150,
3128
},
3229
},
3330
},

internal/graders/behavior_grader.go

Lines changed: 19 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -83,24 +83,27 @@ func (bg *behaviorGrader) Grade(ctx context.Context, gradingContext *Context) (*
8383
feedback = strings.Join(failures, "; ")
8484
}
8585

86+
details := map[string]any{
87+
"max_tool_calls": bg.maxToolCalls,
88+
"max_tokens": bg.maxTokens,
89+
"required_tools": bg.requiredTools,
90+
"forbidden_tools": bg.forbiddenTools,
91+
"max_duration_ms": bg.maxDurationMS,
92+
"failures": failures,
93+
"tool_call_count": session.ToolCallCount,
94+
"tools_used": session.ToolsUsed,
95+
"actual_duration": gradingContext.DurationMS,
96+
}
97+
if session.Usage != nil {
98+
details["tokens_total"] = session.Usage.InputTokens + session.Usage.OutputTokens
99+
}
86100
return &models.GraderResults{
87101
Name: bg.name,
88102
Type: models.GraderKindBehavior,
89103
Score: score,
90104
Passed: len(failures) == 0,
91105
Feedback: feedback,
92-
Details: map[string]any{
93-
"max_tool_calls": bg.maxToolCalls,
94-
"max_tokens": bg.maxTokens,
95-
"required_tools": bg.requiredTools,
96-
"forbidden_tools": bg.forbiddenTools,
97-
"max_duration_ms": bg.maxDurationMS,
98-
"failures": failures,
99-
"tool_call_count": session.ToolCallCount,
100-
"tokens_total": session.TokensTotal,
101-
"tools_used": session.ToolsUsed,
102-
"actual_duration": gradingContext.DurationMS,
103-
},
106+
Details: details,
104107
}, nil
105108
})
106109
}
@@ -123,8 +126,10 @@ func (bg *behaviorGrader) checkMaxTokens(session *models.SessionDigest) []string
123126
return nil
124127
}
125128

126-
if session.TokensTotal > bg.maxTokens {
127-
return []string{fmt.Sprintf("Token usage %d exceeds max allowed %d", session.TokensTotal, bg.maxTokens)}
129+
if session.Usage == nil {
130+
return []string{"No usage data available to check max tokens"}
131+
} else if total := session.Usage.InputTokens + session.Usage.OutputTokens; total > bg.maxTokens {
132+
return []string{fmt.Sprintf("Token usage %d exceeds max allowed %d", total, bg.maxTokens)}
128133
}
129134
return nil
130135
}

internal/graders/behavior_grader_test.go

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,7 @@ func TestBehaviorGrader_MaxToolCalls(t *testing.T) {
9191
results, err := g.Grade(context.Background(), &Context{
9292
Session: &models.SessionDigest{
9393
ToolCallCount: 999,
94-
TokensTotal: 500,
94+
Usage: &models.UsageStats{InputTokens: 500},
9595
},
9696
})
9797
require.NoError(t, err)
@@ -109,7 +109,7 @@ func TestBehaviorGrader_MaxTokens(t *testing.T) {
109109

110110
results, err := g.Grade(context.Background(), &Context{
111111
Session: &models.SessionDigest{
112-
TokensTotal: 500,
112+
Usage: &models.UsageStats{InputTokens: 500},
113113
},
114114
})
115115
require.NoError(t, err)
@@ -125,7 +125,7 @@ func TestBehaviorGrader_MaxTokens(t *testing.T) {
125125

126126
results, err := g.Grade(context.Background(), &Context{
127127
Session: &models.SessionDigest{
128-
TokensTotal: 1000,
128+
Usage: &models.UsageStats{InputTokens: 1000},
129129
},
130130
})
131131
require.NoError(t, err)
@@ -141,7 +141,7 @@ func TestBehaviorGrader_MaxTokens(t *testing.T) {
141141

142142
results, err := g.Grade(context.Background(), &Context{
143143
Session: &models.SessionDigest{
144-
TokensTotal: 1200,
144+
Usage: &models.UsageStats{InputTokens: 1200},
145145
},
146146
})
147147
require.NoError(t, err)
@@ -312,7 +312,7 @@ func TestBehaviorGrader_CombinedRules(t *testing.T) {
312312
DurationMS: 5000,
313313
Session: &models.SessionDigest{
314314
ToolCallCount: 3,
315-
TokensTotal: 800,
315+
Usage: &models.UsageStats{InputTokens: 800},
316316
ToolsUsed: []string{"read_file", "write_file"},
317317
},
318318
})
@@ -335,7 +335,7 @@ func TestBehaviorGrader_CombinedRules(t *testing.T) {
335335
DurationMS: 5000,
336336
Session: &models.SessionDigest{
337337
ToolCallCount: 15, // over limit
338-
TokensTotal: 800,
338+
Usage: &models.UsageStats{InputTokens: 800},
339339
ToolsUsed: []string{"read_file", "write_file"},
340340
},
341341
})
@@ -357,9 +357,9 @@ func TestBehaviorGrader_CombinedRules(t *testing.T) {
357357

358358
results, err := g.Grade(context.Background(), &Context{
359359
Session: &models.SessionDigest{
360-
ToolCallCount: 20, // over limit
361-
TokensTotal: 5000, // over limit
362-
ToolsUsed: []string{"exec"}, // forbidden used, required missing
360+
ToolCallCount: 20, // over limit
361+
Usage: &models.UsageStats{InputTokens: 5000}, // over limit
362+
ToolsUsed: []string{"exec"}, // forbidden used, required missing
363363
},
364364
})
365365
require.NoError(t, err)
@@ -378,7 +378,7 @@ func TestBehaviorGrader_CombinedRules(t *testing.T) {
378378
results, err := g.Grade(context.Background(), &Context{
379379
Session: &models.SessionDigest{
380380
ToolCallCount: 10,
381-
TokensTotal: 500,
381+
Usage: &models.UsageStats{InputTokens: 500},
382382
},
383383
})
384384
require.NoError(t, err)
@@ -430,7 +430,7 @@ func TestBehaviorGrader_EdgeCases(t *testing.T) {
430430

431431
results, err := g.Grade(context.Background(), &Context{
432432
Session: &models.SessionDigest{
433-
TokensTotal: 0,
433+
Usage: &models.UsageStats{InputTokens: 0},
434434
},
435435
})
436436
require.NoError(t, err)

internal/graders/tool_constraint_grader.go

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -136,20 +136,23 @@ func (tc *toolConstraintGrader) Grade(ctx context.Context, gradingContext *Conte
136136
feedback = strings.Join(failures, "; ")
137137
}
138138

139+
details := map[string]any{
140+
"expect_tools": describeToolSpecs(tc.expectTools),
141+
"reject_tools": describeToolSpecs(tc.rejectTools),
142+
"failures": failures,
143+
"tools_used": session.ToolsUsed,
144+
}
145+
if session.Usage != nil {
146+
details["tokens_total"] = session.Usage.InputTokens + session.Usage.OutputTokens
147+
details["total_turns"] = session.Usage.Turns
148+
}
139149
return &models.GraderResults{
140150
Name: tc.name,
141151
Type: models.GraderKindToolConstraint,
142152
Score: score,
143153
Passed: len(failures) == 0,
144154
Feedback: feedback,
145-
Details: map[string]any{
146-
"expect_tools": describeToolSpecs(tc.expectTools),
147-
"reject_tools": describeToolSpecs(tc.rejectTools),
148-
"failures": failures,
149-
"tools_used": session.ToolsUsed,
150-
"total_turns": session.TotalTurns,
151-
"tokens_total": session.TokensTotal,
152-
},
155+
Details: details,
153156
}, nil
154157
})
155158
}

internal/graders/tool_constraint_grader_test.go

Lines changed: 6 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -125,10 +125,9 @@ func TestToolConstraintGrader_AllConstraints_Pass(t *testing.T) {
125125

126126
result, err := g.Grade(context.Background(), &Context{
127127
Session: &models.SessionDigest{
128-
ToolsUsed: []string{"bash", "edit", "view"},
129-
ToolCalls: []models.ToolCall{{Name: "bash"}, {Name: "edit"}, {Name: "view"}},
130-
TotalTurns: 10,
131-
TokensTotal: 4000,
128+
ToolsUsed: []string{"bash", "edit", "view"},
129+
ToolCalls: []models.ToolCall{{Name: "bash"}, {Name: "edit"}, {Name: "view"}},
130+
Usage: &models.UsageStats{Turns: 10, InputTokens: 4000},
132131
},
133132
})
134133
if err != nil {
@@ -152,10 +151,9 @@ func TestToolConstraintGrader_AllConstraints_PartialFail(t *testing.T) {
152151
// bash used, edit missing, create_file used
153152
result, err := g.Grade(context.Background(), &Context{
154153
Session: &models.SessionDigest{
155-
ToolsUsed: []string{"bash", "create_file"},
156-
ToolCalls: []models.ToolCall{{Name: "bash"}, {Name: "create_file"}},
157-
TotalTurns: 10,
158-
TokensTotal: 8000,
154+
ToolsUsed: []string{"bash", "create_file"},
155+
ToolCalls: []models.ToolCall{{Name: "bash"}, {Name: "create_file"}},
156+
Usage: &models.UsageStats{Turns: 10, InputTokens: 8000},
159157
},
160158
})
161159
require.NoError(t, err)

internal/metrics/behavior.go

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,9 +23,13 @@ type BehaviorMetrics struct {
2323
// ComputeBehaviorMetrics analyzes a RunResult against BehaviorRules and returns
2424
// quality metrics including compliance checks and an efficiency score.
2525
func ComputeBehaviorMetrics(run *models.RunResult, rules *models.BehaviorRules) *BehaviorMetrics {
26+
iterationCount := 0
27+
if run.SessionDigest.Usage != nil {
28+
iterationCount = run.SessionDigest.Usage.Turns
29+
}
2630
m := &BehaviorMetrics{
2731
ToolCallCount: run.SessionDigest.ToolCallCount,
28-
IterationCount: run.SessionDigest.TotalTurns,
32+
IterationCount: iterationCount,
2933
}
3034

3135
toolSet := make(map[string]bool, len(run.SessionDigest.ToolsUsed))

0 commit comments

Comments
 (0)