feat: show usage of messages

Author: Neet-Nestor

app/client/api.ts

@@ -1,3 +1,4 @@
+import { CompletionUsage } from "@neet-nestor/web-llm";
 import { CacheType, ModelType } from "../store";
 export const ROLES = ["system", "user", "assistant"] as const;
 export type MessageRole = (typeof ROLES)[number];
@@ -33,7 +34,7 @@ export interface ChatOptions {
   config: LLMConfig;
 
   onUpdate?: (message: string, chunk: string) => void;
-  onFinish: (message: string) => void;
+  onFinish: (message: string, usage?: CompletionUsage) => void;
   onError?: (err: Error) => void;
   onController?: (controller: AbortController) => void;
 }

app/client/webllm.ts

@@ -10,6 +10,7 @@ import {
   ChatCompletionChunk,
   ChatCompletion,
   WebWorkerMLCEngine,
+  CompletionUsage,
 } from "@neet-nestor/web-llm";
 
 import { ChatOptions, LLMApi, LLMConfig, RequestMessage } from "./api";
@@ -107,12 +108,15 @@ export class WebLLMApi implements LLMApi {
     }
 
     let reply: string | null = "";
+    let usage: CompletionUsage | undefined = undefined;
     try {
-      reply = await this.chatCompletion(
+      const completion = await this.chatCompletion(
         !!options.config.stream,
         options.messages,
         options.onUpdate,
       );
+      reply = completion.content;
+      usage = completion.usage;
     } catch (err: any) {
       let errorMessage = err.message || err.toString() || "";
       if (errorMessage === "[object Object]") {
@@ -147,11 +151,13 @@ export class WebLLMApi implements LLMApi {
         return;
       }
       try {
-        reply = await this.chatCompletion(
+        const completion = await this.chatCompletion(
           !!options.config.stream,
           options.messages,
           options.onUpdate,
         );
+        reply = completion.content;
+        usage = completion.usage;
       } catch (err: any) {
         let errorMessage = err.message || err.toString() || "";
         if (errorMessage === "[object Object]") {
@@ -164,7 +170,7 @@ export class WebLLMApi implements LLMApi {
     }
 
     if (reply) {
-      options.onFinish(reply);
+      options.onFinish(reply, usage);
     } else {
       options.onError?.(new Error("Empty response generated by LLM"));
     }
@@ -216,26 +222,38 @@ export class WebLLMApi implements LLMApi {
   async chatCompletion(
     stream: boolean,
     messages: RequestMessage[],
-    onUpdate?: (message: string, chunk: string) => void,
+    onUpdate?: (
+      message: string,
+      chunk: string,
+      usage?: CompletionUsage,
+    ) => void,
   ) {
-    let reply: string | null = "";
-
     const completion = await this.webllm.engine.chatCompletion({
       stream: stream,
       messages: messages as ChatCompletionMessageParam[],
+      stream_options: { include_usage: true },
     });
 
     if (stream) {
+      let content: string | null = "";
+      let usage: CompletionUsage | undefined = undefined;
       const asyncGenerator = completion as AsyncIterable<ChatCompletionChunk>;
       for await (const chunk of asyncGenerator) {
-        if (chunk.choices[0].delta.content) {
-          reply += chunk.choices[0].delta.content;
-          onUpdate?.(reply, chunk.choices[0].delta.content);
+        if (chunk.choices[0]?.delta.content) {
+          content += chunk.choices[0].delta.content;
+          onUpdate?.(content, chunk.choices[0].delta.content);
+        }
+        if (chunk.usage) {
+          usage = chunk.usage;
         }
       }
-      return reply;
+      return { content, usage };
     }
-    return (completion as ChatCompletion).choices[0].message.content;
+
+    return {
+      content: (completion as ChatCompletion).choices[0].message.content,
+      usage: (completion as ChatCompletion).usage,
+    };
   }
 }
 

app/components/chat.module.scss

@@ -518,8 +518,16 @@
   width: 100%;
   box-sizing: border-box;
   padding-right: 10px;
-  pointer-events: none;
   z-index: 1;
+  display: flex;
+  justify-content: end;
+  align-items: center;
+  gap: 4px;
+
+  .tooltip {
+    width: 14px;
+    height: 14px;
+  }
 }
 
 .chat-message-user>.chat-message-container>.chat-message-item {

app/components/chat.tsx

@@ -488,7 +488,6 @@ export function ChatActions(props: {
   showPromptHints: () => void;
   hitBottom: boolean;
   uploading: boolean;
-  tootip?: ReactElement;
 }) {
   const config = useAppConfig();
   const chatStore = useChatStore();
@@ -560,19 +559,6 @@ export function ChatActions(props: {
         icon={<RobotIcon />}
         fullWidth
       />
-      {props.tootip && (
-        <div className={styles.tooltip}>
-          <Tooltip
-            direction="left"
-            content={
-              <div style={{ fontSize: config.fontSize }}>{props.tootip}</div>
-            }
-          >
-            {<InfoIcon />}
-          </Tooltip>
-        </div>
-      )}
-
       {showModelSelector && (
         <Selector
           defaultSelectedValue={currentModel}
@@ -702,23 +688,6 @@ function _Chat() {
     }
   };
 
-  const getEngineStats = () => {
-    webllm.webllm.engine.runtimeStatsText().then((stats) => {
-      const lines = stats.split(", ");
-      setEngineStats(
-        <>
-          <b>WebLLM Engine Statistics</b>
-          {lines.map((line) => (
-            <>
-              <br />
-              <span>{line}</span>
-            </>
-          ))}
-        </>,
-      );
-    });
-  };
-
   const onSubmit = (userInput: string) => {
     if (userInput.trim() === "") return;
 
@@ -732,7 +701,7 @@ function _Chat() {
 
     if (isStreaming) return;
 
-    chatStore.onUserInput(userInput, webllm, attachImages, getEngineStats);
+    chatStore.onUserInput(userInput, webllm, attachImages);
     setAttachImages([]);
     localStorage.setItem(LAST_INPUT_KEY, userInput);
     setUserInput("");
@@ -883,7 +852,7 @@ function _Chat() {
     // resend the message
     const textContent = getMessageTextContent(userMessage);
     const images = getMessageImages(userMessage);
-    chatStore.onUserInput(textContent, webllm, images, getEngineStats);
+    chatStore.onUserInput(textContent, webllm, images);
     inputRef.current?.focus();
   };
 
@@ -1366,6 +1335,38 @@ function _Chat() {
                   </div>
 
                   <div className={styles["chat-message-action-date"]}>
+                    {message.usage && (
+                      <div className={styles.tooltip}>
+                        <Tooltip
+                          direction="top"
+                          content={
+                            <div style={{ fontSize: config.fontSize }}>
+                              <span>
+                                {`Prompt Tokens: ${message.usage.prompt_tokens}`}
+                              </span>
+                              <br />
+                              <span>
+                                {`Completion Tokens: ${message.usage.completion_tokens}`}
+                              </span>
+                              <br />
+                              <span>
+                                {`Prefill: ${message.usage.extra.prefill_tokens_per_s.toFixed(
+                                  4,
+                                )} tokens/sec`}
+                              </span>
+                              <br />
+                              <span>
+                                {`Decoding: ${message.usage.extra.decode_tokens_per_s.toFixed(
+                                  4,
+                                )} tokens/sec`}
+                              </span>
+                            </div>
+                          }
+                        >
+                          {<InfoIcon />}
+                        </Tooltip>
+                      </div>
+                    )}
                     {isContext
                       ? Locale.Chat.IsContext
                       : message.date.toLocaleString()}
@@ -1388,7 +1389,6 @@ function _Chat() {
           scrollToBottom={scrollToBottom}
           hitBottom={hitBottom}
           uploading={uploading}
-          tootip={engineStat}
           showPromptSetting={() => setShowEditPromptModal(true)}
           showPromptHints={() => {
             // Click again to close

app/components/ui-lib.module.scss

@@ -38,7 +38,20 @@
 
 .tooltip {
   position: relative;
-  display: inline-block;
+  display: flex;
+  align-items: center;
+  height: 100%;
+  width: 100%;
+}
+
+.tooltip-icon {
+  display: flex;
+  align-items: center;
+
+  svg {
+    height: 100%;
+    width: 100%;
+  }
 }
 
 .tooltip-content {
@@ -50,6 +63,7 @@
   z-index: 3;
   white-space: nowrap;
   font-size: 12px;
+  text-align: center;
 
   &.tooltip-content-left {
     right: 100%;

app/components/ui-lib.tsx

@@ -53,7 +53,11 @@ export function Tooltip(props: {
   return (
     <div className={styles.tooltip}>
       {/* The anchor component */}
-      <div onMouseOver={showTooltip} onMouseOut={hideTooltip}>
+      <div
+        className={styles["tooltip-icon"]}
+        onMouseOver={showTooltip}
+        onMouseOut={hideTooltip}
+      >
         {props.children}
       </div>
 

app/icons/info.svg

@@ -12,18 +12,19 @@ import {
   StoreKey,
 } from "../constant";
 import { RequestMessage, MultimodalContent } from "../client/api";
-import { prettyObject } from "../utils/format";
 import { estimateTokenLength } from "../utils/token";
 import { nanoid } from "nanoid";
 import { createPersistStore } from "../utils/store";
 import { WebLLMApi } from "../client/webllm";
+import { CompletionUsage } from "@neet-nestor/web-llm";
 
 export type ChatMessage = RequestMessage & {
   date: string;
   streaming?: boolean;
   isError?: boolean;
   id: string;
   model?: ModelType;
+  usage?: CompletionUsage;
 };
 
 export function createMessage(override: Partial<ChatMessage>): ChatMessage {
@@ -275,12 +276,7 @@ export const useChatStore = createPersistStore(
         get().summarizeSession(webllm);
       },
 
-      onUserInput(
-        content: string,
-        webllm: WebLLMApi,
-        attachImages?: string[],
-        onFinish?: () => void,
-      ) {
+      onUserInput(content: string, webllm: WebLLMApi, attachImages?: string[]) {
         const modelConfig = useAppConfig.getState().modelConfig;
 
         const userContent = fillTemplateWith(content, modelConfig);
@@ -353,16 +349,16 @@ export const useChatStore = createPersistStore(
               session.messages = session.messages.concat();
             });
           },
-          onFinish(message) {
+          onFinish(message, usage) {
             botMessage.streaming = false;
+            botMessage.usage = usage;
             if (message) {
               botMessage.content = message;
               get().onNewMessage(botMessage, webllm);
             }
             get().updateCurrentSession((session) => {
               session.isGenerating = false;
             });
-            onFinish?.();
           },
           onError(error) {
             const errorMessage =