From 7395ed6a76b6d0352052ba7a865c1258c2c24672 Mon Sep 17 00:00:00 2001
From: johnd0e <1838643+johnd0e@users.noreply.github.com>
Date: Fri, 26 Jul 2024 23:36:46 +0200
Subject: [PATCH] Process usage (stream)

https://platform.openai.com/docs/api-reference/chat/streaming#chat/streaming-usage
---
 src/worker.mjs | 30 +++++++++++++++++-------------
 1 file changed, 17 insertions(+), 13 deletions(-)

diff --git a/src/worker.mjs b/src/worker.mjs
index 83c6655..038abcb 100644
--- a/src/worker.mjs
+++ b/src/worker.mjs
@@ -281,11 +281,11 @@ async function parseStreamFlush (controller) {
   }
 }
 
-function transformResponseStream (cand, stop, first) {
-  const item = transformCandidatesDelta(cand);
+function transformResponseStream (data, stop, first) {
+  const item = transformCandidatesDelta(data.candidates[0]);
   if (stop) { item.delta = {}; } else { item.finish_reason = null; }
   if (first) { item.delta.content = ""; } else { delete item.delta.role; }
-  const data = {
+  const output = {
     id: this.id,
     choices: [item],
     created: Math.floor(Date.now()/1000),
@@ -293,40 +293,44 @@ function transformResponseStream (cand, stop, first) {
     //system_fingerprint: "fp_69829325d0",
     object: "chat.completion.chunk",
   };
-  return "data: " + JSON.stringify(data) + delimiter;
+  if (stop && data.usageMetadata) {
+    output.usage = transformUsage(data.usageMetadata);
+  }
+  return "data: " + JSON.stringify(output) + delimiter;
 }
 const delimiter = "\n\n";
 async function toOpenAiStream (chunk, controller) {
   const transform = transformResponseStream.bind(this);
   const line = await chunk;
   if (!line) { return; }
-  let candidates;
+  let data;
   try {
-    candidates = JSON.parse(line).candidates;
+    data = JSON.parse(line);
   } catch (err) {
     console.error(line);
     console.error(err);
     const length = this.last.length || 1; // at least 1 error msg
-    candidates = Array.from({ length }, (_, index) => ({
+    const candidates = Array.from({ length }, (_, index) => ({
       finishReason: "error",
       content: { parts: [{ text: err }] },
       index,
     }));
+    data = { candidates };
   }
-  const cand = candidates[0]; // !!untested with candidateCount>1
+  const cand = data.candidates[0]; // !!untested with candidateCount>1
   if (!this.last[cand.index]) {
-    controller.enqueue(transform(cand, false, "first"));
+    controller.enqueue(transform(data, false, "first"));
   }
-  this.last[cand.index] = cand;
+  this.last[cand.index] = data;
   if (cand.content) { // prevent empty data (e.g. when MAX_TOKENS)
-    controller.enqueue(transform(cand));
+    controller.enqueue(transform(data));
   }
 }
 async function toOpenAiStreamFlush (controller) {
   const transform = transformResponseStream.bind(this);
   if (this.last.length > 0) {
-    for (const cand of this.last) {
-      controller.enqueue(transform(cand, "stop"));
+    for (const data of this.last) {
+      controller.enqueue(transform(data, "stop"));
     }
     controller.enqueue("data: [DONE]" + delimiter);
   }