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
704 completedEvent, _ = sjson.SetBytes(completedEvent, "interaction.id", interactionID)704 completedEvent, _ = sjson.SetBytes(completedEvent, "interaction.id", interactionID)
705 completedEvent, _ = sjson.SetBytes(completedEvent, "interaction.model", req.Model)705 completedEvent, _ = sjson.SetBytes(completedEvent, "interaction.model", req.Model)
706 if finalUsage != nil {706 if finalUsage != nil {
707 completedEvent, _ = sjson.SetBytes(completedEvent, "interaction.usage.total_input_tokens", finalUsage.PromptTokens)707 totalInput := finalUsage.PromptTokens + finalUsage.CachedTokens
708 completedEvent, _ = sjson.SetBytes(completedEvent, "interaction.usage.total_output_tokens", finalUsage.CompletionTokens)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)
709 completedEvent, _ = sjson.SetBytes(completedEvent, "interaction.usage.total_cached_tokens", finalUsage.CachedTokens)713 completedEvent, _ = sjson.SetBytes(completedEvent, "interaction.usage.total_cached_tokens", finalUsage.CachedTokens)
714 completedEvent, _ = sjson.SetBytes(completedEvent, "interaction.usage.total_tokens", totalTokens)
710 if detail, ok := helps.ParseInteractionsStreamUsage(completedEvent); ok {715 if detail, ok := helps.ParseInteractionsStreamUsage(completedEvent); ok {
711 reporter.Publish(ctx, detail)716 reporter.Publish(ctx, detail)
712 }717 }
@@ -895,9 +900,14 @@ func consumeDevinFramesToInteractions(body io.Reader, model, chatModelUID string · 182 unmodified lines
895 }900 }
896 901
897 if finalUsage != nil {902 if finalUsage != nil {
898 out, _ = sjson.SetBytes(out, "usage.total_input_tokens", finalUsage.PromptTokens)903 totalInput := finalUsage.PromptTokens + finalUsage.CachedTokens
899 out, _ = sjson.SetBytes(out, "usage.total_output_tokens", finalUsage.CompletionTokens)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)
900 out, _ = sjson.SetBytes(out, "usage.total_cached_tokens", finalUsage.CachedTokens)909 out, _ = sjson.SetBytes(out, "usage.total_cached_tokens", finalUsage.CachedTokens)
910 out, _ = sjson.SetBytes(out, "usage.total_tokens", totalTokens)
901 }911 }
902 912
903 respLog := &helps.DevinUpstreamResponseLog{913 respLog := &helps.DevinUpstreamResponseLog{
internal/runtime/executor/devin_executor_test.go+5−2
@@ -351,8 +351,8 @@ func TestConsumeDevinFramesToInteractions(t *testing.T) { · 350 unmodified lines
351 if root.Get("status").String() != "completed" {351 if root.Get("status").String() != "completed" {
352 t.Errorf("status = %q, want completed", root.Get("status").String())352 t.Errorf("status = %q, want completed", root.Get("status").String())
353 }353 }
354 if root.Get("usage.total_input_tokens").Int() != 100 {354 if root.Get("usage.total_input_tokens").Int() != 120 {
355 t.Errorf("input tokens = %d, want 100", root.Get("usage.total_input_tokens").Int())355 t.Errorf("input tokens = %d, want 120", root.Get("usage.total_input_tokens").Int())
356 }356 }
357 if root.Get("usage.total_output_tokens").Int() != 50 {357 if root.Get("usage.total_output_tokens").Int() != 50 {
358 t.Errorf("output tokens = %d, want 50", root.Get("usage.total_output_tokens").Int())358 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
360 if root.Get("usage.total_cached_tokens").Int() != 20 {360 if root.Get("usage.total_cached_tokens").Int() != 20 {
361 t.Errorf("cached tokens = %d, want 20", root.Get("usage.total_cached_tokens").Int())361 t.Errorf("cached tokens = %d, want 20", root.Get("usage.total_cached_tokens").Int())
362 }362 }
363 if root.Get("usage.total_tokens").Int() != 170 {
364 t.Errorf("total tokens = %d, want 170", root.Get("usage.total_tokens").Int())
365 }
363 366
364 steps := root.Get("steps").Array()367 steps := root.Get("steps").Array()
365 if len(steps) != 3 {368 if len(steps) != 3 {