chore: update branch with main

CZZ000 · Mar 9, 2024 · 3439bcb · 3439bcb
1 parent fc98459
commit 3439bcb
Showing 1 changed file with 11 additions and 7 deletions.
diff --git a/src/providers/openai/chatComplete.ts b/src/providers/openai/chatComplete.ts
@@ -99,19 +99,23 @@ export const OpenAIChatCompleteResponseTransform: (response: OpenAIChatCompleteR
 export const OpenAIChatCompleteJSONToStreamResponseTransform: (response: OpenAIChatCompleteResponse, provider: string) => Array<string> = (response, provider) => {
   const streamChunkArray: Array<string> = [];
   const { id, model, system_fingerprint, choices } = response;
-  const streamChunkTemplate = {
+
+  const { prompt_tokens, completion_tokens } = response.usage || {};
+
+  let total_tokens;
+  if (prompt_tokens && completion_tokens) total_tokens = prompt_tokens + completion_tokens;
+
+  const streamChunkTemplate: Record<string, any> = {
     id,
     object: "chat.completion.chunk",
     created: Date.now(),
     model: model || "",
     system_fingerprint: system_fingerprint || null,
     provider,
-    usage: {}
-  }
-
-  if (response.usage?.completion_tokens) {
-    streamChunkTemplate.usage = {
-      completion_tokens: response.usage?.completion_tokens
+    usage: {
+      ...(completion_tokens && {completion_tokens}),
+      ...(prompt_tokens && {prompt_tokens}),
+      ...(total_tokens && {total_tokens})
     }
   }