Recalculate Tokens Per Second (TPS) using a rolling average based on the time and tokens since the last report to ensure accurate streaming metrics

Author: robertpiosik

packages/vscode/src/utils/make-api-request.ts

@@ -75,28 +75,41 @@ export async function make_api_request(params: {
 
   let stream_start_time = 0
   let last_on_chunk_time = 0
+  let last_total_tokens = 0
   let total_tokens = 0
+  let first_chunk_reported = false
 
   const handle_chunk_metrics = (chunk: string) => {
     if (!params.on_chunk || !chunk) return
 
     total_tokens += Math.ceil(chunk.length / 4)
-    if (stream_start_time == 0) {
+    if (stream_start_time === 0) {
       stream_start_time = Date.now()
       last_on_chunk_time = stream_start_time
       // Don't calculate TPS on the first chunk, as elapsed time is 0
       return
     }
 
     const current_time = Date.now()
-    if (current_time - last_on_chunk_time >= 1000) {
-      last_on_chunk_time = current_time
+    const elapsed_since_last = current_time - last_on_chunk_time
+    if (elapsed_since_last >= 1000) {
+      // Skip the very first report to avoid inflated TPS
+      if (!first_chunk_reported) {
+        first_chunk_reported = true
+        last_on_chunk_time = current_time
+        last_total_tokens = total_tokens
+        return
+      }
 
-      const elapsed_seconds = (current_time - stream_start_time) / 1000
+      const tokens_since_last = total_tokens - last_total_tokens
+      const seconds_since_last = elapsed_since_last / 1000
       const tokens_per_second =
-        elapsed_seconds > 0 ? Math.round(total_tokens / elapsed_seconds) : 0
-
+        seconds_since_last > 0
+          ? Math.round(tokens_since_last / seconds_since_last)
+          : 0
       params.on_chunk(tokens_per_second, total_tokens)
+      last_on_chunk_time = current_time
+      last_total_tokens = total_tokens
     }
   }