💄 style: show token generate performance

Author: cy948

src/features/Conversation/Extras/Usage/UsageDetail/index.tsx

@@ -9,27 +9,27 @@ import { Center, Flexbox } from 'react-layout-kit';
 import { aiModelSelectors, useAiInfraStore } from '@/store/aiInfra';
 import { useGlobalStore } from '@/store/global';
 import { systemStatusSelectors } from '@/store/global/selectors';
-import { ModelTokensUsage } from '@/types/message';
+import { MessageMetadata } from '@/types/message';
 import { formatNumber } from '@/utils/format';
 
 import ModelCard from './ModelCard';
 import TokenProgress, { TokenProgressItem } from './TokenProgress';
 import { getDetailsToken } from './tokens';
 
 interface TokenDetailProps {
+  meta: MessageMetadata;
   model: string;
   provider: string;
-  usage: ModelTokensUsage;
 }
 
-const TokenDetail = memo<TokenDetailProps>(({ usage, model, provider }) => {
+const TokenDetail = memo<TokenDetailProps>(({ meta, model, provider }) => {
   const { t } = useTranslation('chat');
   const theme = useTheme();
 
   const modelCard = useAiInfraStore(aiModelSelectors.getModelCard(model, provider));
   const isShowCredit = useGlobalStore(systemStatusSelectors.isShowCredit) && !!modelCard?.pricing;
 
-  const detailTokens = getDetailsToken(usage, modelCard);
+  const detailTokens = getDetailsToken(meta, modelCard);
   const inputDetails = [
     !!detailTokens.inputAudio && {
       color: theme.cyan9,
@@ -113,6 +113,10 @@ const TokenDetail = memo<TokenDetailProps>(({ usage, model, provider }) => {
     detailTokens.totalTokens!.credit / detailTokens.totalTokens!.token,
     2,
   );
+
+  const tps = meta?.tps ? formatNumber(meta.tps, 2) : undefined;
+  const ttft = meta?.ttft ? formatNumber(meta.ttft / 1000, 2) : undefined;
+
   return (
     <Popover
       arrow={false}
@@ -170,6 +174,18 @@ const TokenDetail = memo<TokenDetailProps>(({ usage, model, provider }) => {
                   <div style={{ fontWeight: 500 }}>{averagePricing}</div>
                 </Flexbox>
               )}
+              {tps && (
+                <Flexbox align={'center'} gap={4} horizontal justify={'space-between'}>
+                  <div style={{ color: theme.colorTextSecondary }}>{'Token Speed (t/s)'}</div>
+                  <div style={{ fontWeight: 500 }}>{tps}</div>
+                </Flexbox>
+              )}
+              {ttft && (
+                <Flexbox align={'center'} gap={4} horizontal justify={'space-between'}>
+                  <div style={{ color: theme.colorTextSecondary }}>{'Delay (s)'}</div>
+                  <div style={{ fontWeight: 500 }}>{ttft}</div>
+                </Flexbox>
+              )}
             </Flexbox>
           </Flexbox>
         </Flexbox>

src/features/Conversation/Extras/Usage/index.tsx

@@ -31,7 +31,7 @@ const Usage = memo<UsageProps>(({ model, metadata, provider }) => {
       </Center>
 
       {!!metadata.totalTokens && (
-        <TokenDetail model={model as string} provider={provider} usage={metadata} />
+        <TokenDetail meta={metadata} model={model as string} provider={provider} />
       )}
     </Flexbox>
   );

src/libs/agent-runtime/utils/openaiCompatibleFactory/index.ts

@@ -201,6 +201,7 @@ export const LobeOpenAICompatibleFactory = <T extends Record<string, any> = any>
 
     async chat({ responseMode, ...payload }: ChatStreamPayload, options?: ChatCompetitionOptions) {
       try {
+        const inputStartAt = Date.now();
         const postPayload = chatCompletion?.handlePayload
           ? chatCompletion.handlePayload(payload, this._options)
           : ({
@@ -253,10 +254,14 @@ export const LobeOpenAICompatibleFactory = <T extends Record<string, any> = any>
             debugStream(useForDebugStream).catch(console.error);
           }
 
-          const streamHandler = chatCompletion?.handleStream || OpenAIStream;
-          return StreamingResponse(streamHandler(prod, streamOptions), {
-            headers: options?.headers,
-          });
+          return StreamingResponse(
+            chatCompletion?.handleStream
+              ? chatCompletion.handleStream(prod, streamOptions.callbacks)
+              : OpenAIStream(prod, streamOptions, { inputStartAt }),
+            {
+              headers: options?.headers,
+            },
+          );
         }
 
         if (debug?.chatCompletion?.()) {
@@ -269,10 +274,14 @@ export const LobeOpenAICompatibleFactory = <T extends Record<string, any> = any>
           chatCompletion?.handleTransformResponseToStream || transformResponseToStream;
         const stream = transformHandler(response as unknown as OpenAI.ChatCompletion);
 
-        const streamHandler = chatCompletion?.handleStream || OpenAIStream;
-        return StreamingResponse(streamHandler(stream, streamOptions), {
-          headers: options?.headers,
-        });
+        return StreamingResponse(
+          chatCompletion?.handleStream
+            ? chatCompletion.handleStream(stream, streamOptions.callbacks)
+            : OpenAIStream(stream, streamOptions, { inputStartAt }),
+          {
+            headers: options?.headers,
+          },
+        );
       } catch (error) {
         throw this.handleError(error);
       }

src/libs/agent-runtime/utils/streams/openai.ts

@@ -16,6 +16,7 @@ import {
   createCallbacksTransformer,
   createFirstErrorHandleTransformer,
   createSSEProtocolTransformer,
+  createTokenSpeedCalculator,
   generateToolCallId,
 } from './protocol';
 
@@ -39,6 +40,11 @@ export const transformOpenAIStream = (
   }
 
   try {
+    // performance monitor return
+    if (chunk.id === 'speed') {
+      return { data: chunk.object, id: chunk.id, type: 'speed' };
+    }
+
     // maybe need another structure to add support for multiple choices
     const item = chunk.choices[0];
     if (!item) {
@@ -221,9 +227,14 @@ export interface OpenAIStreamOptions {
   provider?: string;
 }
 
+export interface TraceOptions {
+  inputStartAt?: number;
+}
+
 export const OpenAIStream = (
   stream: Stream<OpenAI.ChatCompletionChunk> | ReadableStream,
   { callbacks, provider, bizErrorTypeTransformer }: OpenAIStreamOptions = {},
+  { inputStartAt }: TraceOptions = {},
 ) => {
   const streamStack: StreamContext = { id: '' };
 
@@ -236,7 +247,8 @@ export const OpenAIStream = (
       // provider like huggingface or minimax will return error in the stream,
       // so in the first Transformer, we need to handle the error
       .pipeThrough(createFirstErrorHandleTransformer(bizErrorTypeTransformer, provider))
-      .pipeThrough(createSSEProtocolTransformer(transformOpenAIStream, streamStack))
+      .pipeThrough(createTokenSpeedCalculator(transformOpenAIStream, { inputStartAt, streamStack }))
+      .pipeThrough(createSSEProtocolTransformer((c) => c, streamStack))
       .pipeThrough(createCallbacksTransformer(callbacks))
   );
 };

src/libs/agent-runtime/utils/streams/protocol.ts

@@ -1,5 +1,5 @@
 import { ChatStreamCallbacks } from '@/libs/agent-runtime';
-import { ModelTokensUsage } from '@/types/message';
+import { ModelSpeed, ModelTokensUsage } from '@/types/message';
 
 import { AgentRuntimeErrorType } from '../../error';
 
@@ -50,6 +50,8 @@ export interface StreamProtocolChunk {
     | 'error'
     // token usage
     | 'usage'
+    // performance monitor
+    | 'speed'
     // unknown data result
     | 'data';
 }
@@ -244,3 +246,44 @@ export const createSSEDataExtractor = () =>
       }
     },
   });
+
+/**
+ * Create a middleware to calculate the token generate speed
+ * @requires createSSEProtocolTransformer
+ */
+export const createTokenSpeedCalculator = (
+  transformer: (chunk: any, stack: StreamContext) => StreamProtocolChunk | StreamProtocolChunk[],
+  { streamStack, inputStartAt }: { inputStartAt?: number; streamStack?: StreamContext } = {},
+) => {
+  let outputStartAt: number | undefined;
+
+  const process = (chunk: StreamProtocolChunk) => {
+    let result = [chunk];
+    // if the chunk is the first text chunk, set as output start
+    if (!outputStartAt && chunk.type === 'text') outputStartAt = Date.now();
+    // if the chunk is the stop chunk, set as output finish
+    if (inputStartAt && outputStartAt && chunk.type === 'usage') {
+      const outputTokens = chunk.data?.totalOutputTokens || chunk.data?.outputTextTokens;
+      result.push({
+        data: {
+          tps: (outputTokens / (Date.now() - outputStartAt)) * 1000,
+          ttft: outputStartAt - inputStartAt,
+        } as ModelSpeed,
+        id: 'output_speed',
+        type: 'speed',
+      });
+    }
+    return result;
+  };
+
+  return new TransformStream({
+    transform(chunk, controller) {
+      let result = transformer(chunk, streamStack || { id: '' });
+      if (!Array.isArray(result)) result = [result];
+      result.forEach((r) => {
+        const processed = process(r);
+        if (processed) processed.forEach((p) => controller.enqueue(p));
+      });
+    },
+  });
+};

src/store/chat/slices/aiChat/actions/generateAIChat.ts

@@ -562,7 +562,7 @@ export const generateAIChat: StateCreator<
       },
       onFinish: async (
         content,
-        { traceId, observationId, toolCalls, reasoning, grounding, usage },
+        { traceId, observationId, toolCalls, reasoning, grounding, usage, speed },
       ) => {
         // if there is traceId, update it
         if (traceId) {
@@ -592,13 +592,15 @@ export const generateAIChat: StateCreator<
           internal_toggleToolCallingStreaming(messageId, undefined);
         }
 
+        console.log('Generate Message Finish:', usage, speed);
+
         // update the content after fetch result
         await internal_updateMessageContent(messageId, content, {
           toolCalls,
           reasoning: !!reasoning ? { ...reasoning, duration } : undefined,
           search: !!grounding?.citations ? grounding : undefined,
-          metadata: usage,
           imageList: finalImages.length > 0 ? finalImages : undefined,
+          metadata: speed ? { ...usage, ...speed } : usage,
         });
       },
       onMessageHandle: async (chunk) => {

src/types/message/base.ts

@@ -44,8 +44,16 @@ export interface ModelTokensUsage {
   totalTokens?: number;
 }
 
+export interface ModelSpeed {
+  // tokens per second
+  tps?: number;
+  // time to fist token
+  ttft?: number;
+}
+
 export interface MessageMetadata extends ModelTokensUsage {
   tps?: number;
+  ttft?: number;
 }
 
 export type MessageRoleType = 'user' | 'system' | 'assistant' | 'tool';

src/utils/fetch/fetchSSE.ts

@@ -10,6 +10,7 @@ import {
   MessageToolCallChunk,
   MessageToolCallSchema,
   ModelReasoning,
+  ModelSpeed,
   ModelTokensUsage,
 } from '@/types/message';
 import { ChatImageChunk } from '@/types/message/image';
@@ -29,6 +30,7 @@ export type OnFinishHandler = (
     images?: ChatImageChunk[];
     observationId?: string | null;
     reasoning?: ModelReasoning;
+    speed?: ModelSpeed;
     toolCalls?: MessageToolCall[];
     traceId?: string | null;
     type?: SSEFinishType;
@@ -41,6 +43,11 @@ export interface MessageUsageChunk {
   usage: ModelTokensUsage;
 }
 
+export interface MessageSpeedChunk {
+  speed: ModelSpeed;
+  type: 'speed';
+}
+
 export interface MessageTextChunk {
   text: string;
   type: 'text';
@@ -82,7 +89,8 @@ export interface FetchSSEOptions {
       | MessageReasoningChunk
       | MessageGroundingChunk
       | MessageUsageChunk
-      | MessageBase64ImageChunk,
+      | MessageBase64ImageChunk
+      | MessageSpeedChunk,
   ) => void;
   smoothing?: SmoothingParams | boolean;
 }
@@ -342,6 +350,7 @@ export const fetchSSE = async (url: string, options: RequestInit & FetchSSEOptio
   let grounding: GroundingSearch | undefined = undefined;
   let usage: ModelTokensUsage | undefined = undefined;
   let images: ChatImageChunk[] = [];
+  let speed: ModelSpeed | undefined = undefined;
 
   await fetchEventSource(url, {
     body: options.body,
@@ -433,6 +442,12 @@ export const fetchSSE = async (url: string, options: RequestInit & FetchSSEOptio
           break;
         }
 
+        case 'speed': {
+          speed = data;
+          options.onMessageHandle?.({ speed: data, type: 'speed' });
+          break;
+        }
+
         case 'grounding': {
           grounding = data;
           options.onMessageHandle?.({ grounding: data, type: 'grounding' });
@@ -517,6 +532,7 @@ export const fetchSSE = async (url: string, options: RequestInit & FetchSSEOptio
         images: images.length > 0 ? images : undefined,
         observationId,
         reasoning: !!thinking ? { content: thinking, signature: thinkingSignature } : undefined,
+        speed,
         toolCalls,
         traceId,
         type: finishedType,