Skip to content

Commit 4964e02

Browse files
dvyukovdvyukov
committed
pkg/aflow/trajectory: add token usage
1 parent dd7d1ce commit 4964e02

File tree

8 files changed

+115
-62
lines changed

8 files changed

+115
-62
lines changed

‎dashboard/app/ai.go‎

Lines changed: 41 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -64,20 +64,23 @@ type uiAIResult struct {
6464
}
6565

6666
type uiAITrajectorySpan struct {
67-
Started time.Time
68-
Seq int64
69-
Nesting int64
70-
Type string
71-
Name string
72-
Model string
73-
Duration time.Duration
74-
Error string
75-
Args string
76-
Results string
77-
Instruction string
78-
Prompt string
79-
Reply string
80-
Thoughts string
67+
Started time.Time
68+
Seq int64
69+
Nesting int64
70+
Type string
71+
Name string
72+
Model string
73+
Duration time.Duration
74+
Error string
75+
Args string
76+
Results string
77+
Instruction string
78+
Prompt string
79+
Reply string
80+
Thoughts string
81+
InputTokens int
82+
OutputTokens int
83+
OutputThoughtsTokens int
8184
}
8285

8386
func handleAIJobsPage(ctx context.Context, w http.ResponseWriter, r *http.Request) error {
@@ -233,20 +236,23 @@ func makeUIAITrajectory(trajetory []*aidb.TrajectorySpan) []*uiAITrajectorySpan
233236
duration = span.Finished.Time.Sub(span.Started)
234237
}
235238
res = append(res, &uiAITrajectorySpan{
236-
Started: span.Started,
237-
Seq: span.Seq,
238-
Nesting: span.Nesting,
239-
Type: span.Type,
240-
Name: span.Name,
241-
Model: span.Model,
242-
Duration: duration,
243-
Error: nullString(span.Error),
244-
Args: nullJSON(span.Args),
245-
Results: nullJSON(span.Results),
246-
Instruction: nullString(span.Instruction),
247-
Prompt: nullString(span.Prompt),
248-
Reply: nullString(span.Reply),
249-
Thoughts: nullString(span.Thoughts),
239+
Started: span.Started,
240+
Seq: span.Seq,
241+
Nesting: span.Nesting,
242+
Type: span.Type,
243+
Name: span.Name,
244+
Model: span.Model,
245+
Duration: duration,
246+
Error: nullString(span.Error),
247+
Args: nullJSON(span.Args),
248+
Results: nullJSON(span.Results),
249+
Instruction: nullString(span.Instruction),
250+
Prompt: nullString(span.Prompt),
251+
Reply: nullString(span.Reply),
252+
Thoughts: nullString(span.Thoughts),
253+
InputTokens: nullInt64(span.InputTokens),
254+
OutputTokens: nullInt64(span.OutputTokens),
255+
OutputThoughtsTokens: nullInt64(span.OutputThoughtsTokens),
250256
})
251257
}
252258
return res
@@ -642,3 +648,10 @@ func nullJSON(v spanner.NullJSON) string {
642648
}
643649
return fmt.Sprint(v.Value)
644650
}
651+
652+
func nullInt64(v spanner.NullInt64) int {
653+
if !v.Valid {
654+
return 0
655+
}
656+
return int(v.Int64)
657+
}

‎dashboard/app/aidb/crud.go‎

Lines changed: 25 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -173,21 +173,24 @@ func StoreTrajectorySpan(ctx context.Context, jobID string, span *trajectory.Spa
173173
}
174174
defer client.Close()
175175
ent := TrajectorySpan{
176-
JobID: jobID,
177-
Seq: int64(span.Seq),
178-
Nesting: int64(span.Nesting),
179-
Type: string(span.Type),
180-
Name: span.Name,
181-
Model: span.Model,
182-
Started: span.Started,
183-
Finished: toNullTime(span.Finished),
184-
Error: toNullString(span.Error),
185-
Args: toNullJSON(span.Args),
186-
Results: toNullJSON(span.Results),
187-
Instruction: toNullString(span.Instruction),
188-
Prompt: toNullString(span.Prompt),
189-
Reply: toNullString(span.Reply),
190-
Thoughts: toNullString(span.Thoughts),
176+
JobID: jobID,
177+
Seq: int64(span.Seq),
178+
Nesting: int64(span.Nesting),
179+
Type: string(span.Type),
180+
Name: span.Name,
181+
Model: span.Model,
182+
Started: span.Started,
183+
Finished: toNullTime(span.Finished),
184+
Error: toNullString(span.Error),
185+
Args: toNullJSON(span.Args),
186+
Results: toNullJSON(span.Results),
187+
Instruction: toNullString(span.Instruction),
188+
Prompt: toNullString(span.Prompt),
189+
Reply: toNullString(span.Reply),
190+
Thoughts: toNullString(span.Thoughts),
191+
InputTokens: toNullInt64(span.InputTokens),
192+
OutputTokens: toNullInt64(span.OutputTokens),
193+
OutputThoughtsTokens: toNullInt64(span.OutputThoughtsTokens),
191194
}
192195
mut, err := spanner.InsertOrUpdateStruct("TrajectorySpans", ent)
193196
if err != nil {
@@ -290,3 +293,10 @@ func toNullString(v string) spanner.NullString {
290293
}
291294
return spanner.NullString{StringVal: v, Valid: true}
292295
}
296+
297+
func toNullInt64(v int) spanner.NullInt64 {
298+
if v == 0 {
299+
return spanner.NullInt64{}
300+
}
301+
return spanner.NullInt64{Int64: int64(v), Valid: true}
302+
}

‎dashboard/app/aidb/entities.go‎

Lines changed: 17 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -38,18 +38,21 @@ type Job struct {
3838
type TrajectorySpan struct {
3939
JobID string
4040
// The following fields correspond one-to-one to trajectory.Span fields (add field comments there).
41-
Seq int64
42-
Nesting int64
43-
Type string
44-
Name string
45-
Model string
46-
Started time.Time
47-
Finished spanner.NullTime
48-
Error spanner.NullString
49-
Args spanner.NullJSON
50-
Results spanner.NullJSON
51-
Instruction spanner.NullString
52-
Prompt spanner.NullString
53-
Reply spanner.NullString
54-
Thoughts spanner.NullString
41+
Seq int64
42+
Nesting int64
43+
Type string
44+
Name string
45+
Model string
46+
Started time.Time
47+
Finished spanner.NullTime
48+
Error spanner.NullString
49+
Args spanner.NullJSON
50+
Results spanner.NullJSON
51+
Instruction spanner.NullString
52+
Prompt spanner.NullString
53+
Reply spanner.NullString
54+
Thoughts spanner.NullString
55+
InputTokens spanner.NullInt64
56+
OutputTokens spanner.NullInt64
57+
OutputThoughtsTokens spanner.NullInt64
5558
}

‎dashboard/app/aidb/migrations/5_add_trajectory_tokens.down.sql‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
ALTER TABLE TrajectorySpans DROP COLUMN InputTokens;
2+
ALTER TABLE TrajectorySpans DROP COLUMN OutputTokens;
3+
ALTER TABLE TrajectorySpans DROP COLUMN OutputThoughtsTokens;

‎dashboard/app/aidb/migrations/5_add_trajectory_tokens.up.sql‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,3 @@
1+
ALTER TABLE TrajectorySpans ADD COLUMN InputTokens INT64;
2+
ALTER TABLE TrajectorySpans ADD COLUMN OutputTokens INT64;
3+
ALTER TABLE TrajectorySpans ADD COLUMN OutputThoughtsTokens INT64;

‎dashboard/app/templates/ai_job.html‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,13 @@
8686
{{if $span.Reply}}
8787
<b>Reply:</b> <div id="ai_details_div"><pre>{{$span.Reply}}</pre></div><br>
8888
{{end}}
89+
{{if $span.InputTokens}}
90+
<b>Tokens:</b> <div id="ai_details_div"><pre>
91+
input: {{$span.InputTokens}}
92+
output: {{$span.OutputTokens}}
93+
thoughts: {{$span.OutputThoughtsTokens}}
94+
</pre></div><br>
95+
{{end}}
8996
{{if $span.Thoughts}}
9097
<b>Thoughts:</b> <div id="ai_details_div"><pre>{{$span.Thoughts}}</pre></div><br>
9198
{{end}}

‎pkg/aflow/llm_agent.go‎

Lines changed: 11 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -192,8 +192,7 @@ func (a *LLMAgent) chat(ctx *Context, cfg *genai.GenerateContentConfig, tools ma
192192
if err != nil {
193193
return "", nil, ctx.finishSpan(reqSpan, err)
194194
}
195-
reply, thoughts, calls, respErr := a.parseResponse(resp)
196-
reqSpan.Thoughts = thoughts
195+
reply, calls, respErr := a.parseResponse(resp, reqSpan)
197196
if err := ctx.finishSpan(reqSpan, respErr); err != nil {
198197
return "", nil, err
199198
}
@@ -302,8 +301,8 @@ func (a *LLMAgent) callTools(ctx *Context, tools map[string]Tool, calls []*genai
302301
return responses, outputs, nil
303302
}
304303

305-
func (a *LLMAgent) parseResponse(resp *genai.GenerateContentResponse) (
306-
reply, thoughts string, calls []*genai.FunctionCall, err error) {
304+
func (a *LLMAgent) parseResponse(resp *genai.GenerateContentResponse, span *trajectory.Span) (
305+
reply string, calls []*genai.FunctionCall, err error) {
307306
if len(resp.Candidates) == 0 || resp.Candidates[0] == nil {
308307
err = fmt.Errorf("empty model response")
309308
if resp.PromptFeedback != nil {
@@ -322,6 +321,13 @@ func (a *LLMAgent) parseResponse(resp *genai.GenerateContentResponse) (
322321
err = fmt.Errorf("unexpected reply fields (%+v)", *candidate)
323322
return
324323
}
324+
if resp.UsageMetadata != nil {
325+
// We add ToolUsePromptTokenCount just in case, but Gemini does not use/set it.
326+
span.InputTokens = int(resp.UsageMetadata.PromptTokenCount) +
327+
int(resp.UsageMetadata.ToolUsePromptTokenCount)
328+
span.OutputTokens = int(resp.UsageMetadata.CandidatesTokenCount)
329+
span.OutputThoughtsTokens = int(resp.UsageMetadata.ThoughtsTokenCount)
330+
}
325331
for _, part := range candidate.Content.Parts {
326332
// We don't expect to receive these now.
327333
if part.VideoMetadata != nil || part.InlineData != nil ||
@@ -333,7 +339,7 @@ func (a *LLMAgent) parseResponse(resp *genai.GenerateContentResponse) (
333339
if part.FunctionCall != nil {
334340
calls = append(calls, part.FunctionCall)
335341
} else if part.Thought {
336-
thoughts += part.Text
342+
span.Thoughts += part.Text
337343
} else {
338344
reply += part.Text
339345
}

‎pkg/aflow/trajectory/trajectory.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,12 @@ type Span struct {
3838

3939
// LLM invocation.
4040
Thoughts string `json:",omitzero"`
41+
42+
// For details see:
43+
// https://pkg.go.dev/google.golang.org/genai#GenerateContentResponseUsageMetadata
44+
InputTokens int `json:",omitzero"`
45+
OutputTokens int `json:",omitzero"`
46+
OutputThoughtsTokens int `json:",omitzero"`
4147
}
4248

4349
type SpanType string
@@ -89,6 +95,8 @@ func (span *Span) String() string {
8995
}
9096
fmt.Fprintf(sb, "reply:\n%v\n", span.Reply)
9197
case SpanLLM:
98+
fmt.Fprintf(sb, "tokens: input=%v output=%v thoughts=%v\n",
99+
span.InputTokens, span.OutputTokens, span.OutputThoughtsTokens)
92100
if span.Thoughts != "" {
93101
fmt.Fprintf(sb, "thoughts:\n%v\n", span.Thoughts)
94102
}

0 commit comments

Comments
 (0)