From 87bc89e68388a6ebbfde79f58a3594b849fec52b Mon Sep 17 00:00:00 2001
From: Nestor Qin <imba.qxy@gmail.com>
Date: Sat, 22 Jun 2024 23:14:13 -0400
Subject: [PATCH] feat: collect usage for mlcllm response

---
 app/client/mlcllm.ts | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

diff --git a/app/client/mlcllm.ts b/app/client/mlcllm.ts
index af4bb027..f9fb04a3 100644
--- a/app/client/mlcllm.ts
+++ b/app/client/mlcllm.ts
@@ -1,6 +1,6 @@
 import log from "loglevel";
 import { ChatOptions, LLMApi } from "./api";
-import { ChatCompletionFinishReason } from "@mlc-ai/web-llm";
+import { ChatCompletionFinishReason, CompletionUsage } from "@mlc-ai/web-llm";
 
 export class MlcLLMApi implements LLMApi {
   private endpoint: string;
@@ -24,6 +24,7 @@ export class MlcLLMApi implements LLMApi {
 
     let reply: string = "";
     let stopReason: ChatCompletionFinishReason | undefined;
+    let usage: CompletionUsage | undefined;
 
     try {
       const response = await fetch(`${this.endpoint}/v1/chat/completions`, {
@@ -52,6 +53,10 @@ export class MlcLLMApi implements LLMApi {
               if (data.choices[0].finish_reason) {
                 stopReason = data.choices[0].finish_reason;
               }
+
+              if (data.usage) {
+                usage = data.usage;
+              }
             }
           }
 
@@ -60,12 +65,13 @@ export class MlcLLMApi implements LLMApi {
             break;
           }
         }
-        options.onFinish(reply, stopReason);
+        options.onFinish(reply, stopReason, usage);
       } else {
         const data = await response.json();
         options.onFinish(
           data.choices[0].message.content,
-          data.choices[0].finish_reason,
+          data.choices[0].finish_reason || undefined,
+          data.usage || undefined,
         );
       }
     } catch (error: any) {