From 7c0bf84fed46bc8ddcdad5a4dd3195b9d1312bf7 Mon Sep 17 00:00:00 2001
From: Sai Krishna <saikrishna@Sais-MacBook-Air.local>
Date: Wed, 3 Apr 2024 12:01:45 +0530
Subject: [PATCH] chore: sending usage details if present for stream chunks

---
 src/providers/fireworks-ai/chatComplete.ts | 6 ++++++
 src/providers/fireworks-ai/complete.ts     | 1 +
 2 files changed, 7 insertions(+)

diff --git a/src/providers/fireworks-ai/chatComplete.ts b/src/providers/fireworks-ai/chatComplete.ts
index 52b455112..dd5492110 100644
--- 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
index 6c435bb07..92df53ac8 100644
--- 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'
   );
 };