chore: sending usage details if present for stream chunks

Portkey-AI · Apr 3, 2024 · 7c0bf84 · 7c0bf84
1 parent 71df568
commit 7c0bf84
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 0 deletions.
diff --git a/src/providers/fireworks-ai/chatComplete.ts b/src/providers/fireworks-ai/chatComplete.ts
@@ -118,6 +118,11 @@ export interface FireworksAIStreamChunk {
     index: number;
     finish_reason: string | null;
   }[];
+  usage: null | {
+    prompt_tokens: number;
+    completion_tokens: number;
+    total_tokens: number;
+  };
 }
 
 export const FireworksAIErrorResponseTransform: (
@@ -200,6 +205,7 @@ export const FireworksAIChatCompleteStreamChunkTransform: (
           finish_reason: parsedChunk.choices[0].finish_reason,
         },
       ],
+      ...(parsedChunk.usage ? { usage: parsedChunk.usage } : {}),
     })}` + '\n\n'
   );
 };
diff --git a/src/providers/fireworks-ai/complete.ts b/src/providers/fireworks-ai/complete.ts
@@ -153,6 +153,7 @@ export const FireworksAICompleteStreamChunkTransform: (
           finish_reason: parsedChunk.choices[0].finish_reason,
         },
       ],
+      ...(parsedChunk.usage ? { usage: parsedChunk.usage } : {}),
     })}` + '\n\n'
   );
 };