History

fix(devin): calculate total_input_tokens and total_tokens correctly

sususu committed 2 months ago

Commit 85ddf3a1 parent 0aedd052 files · +19 −6Browse files

Unified · Split

2 files
  1. internal/runtime/executor/devin_executor.go +14 −4
  2. internal/runtime/executor/devin_executor_test.go +5 −2
internal/runtime/executor/devin_executor.go+14−4
@@ -704,9 +704,14 @@ func (e *DevinExecutor) streamDevinFrames( · 703 unmodified lines
704704 completedEvent, _ = sjson.SetBytes(completedEvent, "interaction.id", interactionID)
705705 completedEvent, _ = sjson.SetBytes(completedEvent, "interaction.model", req.Model)
706706 if finalUsage != nil {
707− completedEvent, _ = sjson.SetBytes(completedEvent, "interaction.usage.total_input_tokens", finalUsage.PromptTokens)
708− completedEvent, _ = sjson.SetBytes(completedEvent, "interaction.usage.total_output_tokens", finalUsage.CompletionTokens)
707+ totalInput := finalUsage.PromptTokens + finalUsage.CachedTokens
708+ totalOutput := finalUsage.CompletionTokens
709+ totalTokens := totalInput + totalOutput
710+
711+ completedEvent, _ = sjson.SetBytes(completedEvent, "interaction.usage.total_input_tokens", totalInput)
712+ completedEvent, _ = sjson.SetBytes(completedEvent, "interaction.usage.total_output_tokens", totalOutput)
709713 completedEvent, _ = sjson.SetBytes(completedEvent, "interaction.usage.total_cached_tokens", finalUsage.CachedTokens)
714+ completedEvent, _ = sjson.SetBytes(completedEvent, "interaction.usage.total_tokens", totalTokens)
710715 if detail, ok := helps.ParseInteractionsStreamUsage(completedEvent); ok {
711716 reporter.Publish(ctx, detail)
712717 }
@@ -895,9 +900,14 @@ func consumeDevinFramesToInteractions(body io.Reader, model, chatModelUID string · 182 unmodified lines
895900 }
896901
897902 if finalUsage != nil {
898− out, _ = sjson.SetBytes(out, "usage.total_input_tokens", finalUsage.PromptTokens)
899− out, _ = sjson.SetBytes(out, "usage.total_output_tokens", finalUsage.CompletionTokens)
903+ totalInput := finalUsage.PromptTokens + finalUsage.CachedTokens
904+ totalOutput := finalUsage.CompletionTokens
905+ totalTokens := totalInput + totalOutput
906+
907+ out, _ = sjson.SetBytes(out, "usage.total_input_tokens", totalInput)
908+ out, _ = sjson.SetBytes(out, "usage.total_output_tokens", totalOutput)
900909 out, _ = sjson.SetBytes(out, "usage.total_cached_tokens", finalUsage.CachedTokens)
910+ out, _ = sjson.SetBytes(out, "usage.total_tokens", totalTokens)
901911 }
902912
903913 respLog := &helps.DevinUpstreamResponseLog{
internal/runtime/executor/devin_executor_test.go+5−2
@@ -351,8 +351,8 @@ func TestConsumeDevinFramesToInteractions(t *testing.T) { · 350 unmodified lines
351351 if root.Get("status").String() != "completed" {
352352 t.Errorf("status = %q, want completed", root.Get("status").String())
353353 }
354− if root.Get("usage.total_input_tokens").Int() != 100 {
355− t.Errorf("input tokens = %d, want 100", root.Get("usage.total_input_tokens").Int())
354+ if root.Get("usage.total_input_tokens").Int() != 120 {
355+ t.Errorf("input tokens = %d, want 120", root.Get("usage.total_input_tokens").Int())
356356 }
357357 if root.Get("usage.total_output_tokens").Int() != 50 {
358358 t.Errorf("output tokens = %d, want 50", root.Get("usage.total_output_tokens").Int())
@@ -360,6 +360,9 @@ func TestConsumeDevinFramesToInteractions(t *testing.T) { · 1 unmodified line
360360 if root.Get("usage.total_cached_tokens").Int() != 20 {
361361 t.Errorf("cached tokens = %d, want 20", root.Get("usage.total_cached_tokens").Int())
362362 }
363+ if root.Get("usage.total_tokens").Int() != 170 {
364+ t.Errorf("total tokens = %d, want 170", root.Get("usage.total_tokens").Int())
365+ }
363366
364367 steps := root.Get("steps").Array()
365368 if len(steps) != 3 {