Add token usage tracking and UI refinements

This commit is contained in:
temp
2026-08-24 00:05:05 +08:00
parent a34cab431a
commit e218316c55
23 changed files with 1285 additions and 93 deletions
@@ -371,5 +371,9 @@ func (s *einoRunRuntimeSession) emitUsageSummary(reason string) bool {
if s == nil || s.usage == nil {
return false
}
return s.usage.EmitOnce(s.conversationID, s.orchMode, reason, s.progress, s.logger)
modelName := ""
if s.args != nil {
modelName = s.args.ModelName
}
return s.usage.EmitOnce(s.conversationID, s.orchMode, reason, modelName, s.progress, s.logger)
}
@@ -61,6 +61,7 @@ func (a *einoRunUsageAccumulator) EmitOnce(
conversationID string,
orchestration string,
reason string,
modelName string,
progress func(eventType, message string, data interface{}),
logger *zap.Logger,
) bool {
@@ -81,6 +82,7 @@ func (a *einoRunUsageAccumulator) EmitOnce(
"source": "eino",
"orchestration": orchestration,
"reason": reason,
"model": modelName,
"modelCalls": s.ModelCalls,
"promptTokens": s.PromptTokens,
"completionTokens": s.CompletionTokens,
@@ -96,6 +98,7 @@ func (a *einoRunUsageAccumulator) EmitOnce(
zap.String("conversationId", conversationID),
zap.String("orchestration", orchestration),
zap.String("reason", reason),
zap.String("model", modelName),
zap.Int("modelCalls", s.ModelCalls),
zap.Int("promptTokens", s.PromptTokens),
zap.Int("completionTokens", s.CompletionTokens),
@@ -49,16 +49,16 @@ func TestEinoRunUsageAccumulatorEmitOnce(t *testing.T) {
}
}
if !acc.EmitOnce("conv-1", "deep", "final", progress, nil) {
if !acc.EmitOnce("conv-1", "deep", "final", "gpt-test", progress, nil) {
t.Fatal("first emit should return true")
}
if acc.EmitOnce("conv-1", "deep", "partial", progress, nil) {
if acc.EmitOnce("conv-1", "deep", "partial", "gpt-test", progress, nil) {
t.Fatal("second emit should return false")
}
if len(events) != 1 {
t.Fatalf("events = %#v, want one usage summary", events)
}
if events[0]["conversationId"] != "conv-1" || events[0]["orchestration"] != "deep" || events[0]["reason"] != "final" || events[0]["totalTokens"] != 3 {
if events[0]["conversationId"] != "conv-1" || events[0]["orchestration"] != "deep" || events[0]["reason"] != "final" || events[0]["model"] != "gpt-test" || events[0]["totalTokens"] != 3 {
t.Fatalf("event = %#v", events[0])
}
}