From eb1c7bd2fcfb7475fc20042926f92478f708151e Mon Sep 17 00:00:00 2001 From: Xiaohan-Tian Date: Thu, 14 May 2026 19:05:13 -0700 Subject: [PATCH] feat: added performance statistics info when using embedded LLM --- src/agent/core/AgentCore.ts | 5 +++-- src/agent/llm/LocalBrowserLLMProvider.ts | 10 +++++++++- src/agent/llm/StreamingTypes.ts | 8 +++++++- src/components/ChatBox.css | 8 +++++++- src/components/ChatBox.tsx | 1 + src/components/chat/AssistantMessage.tsx | 23 ++++++++++++++++++++++- src/hooks/useStreamProcessor.ts | 9 +++++++-- src/types/projectTypes.ts | 2 ++ 8 files changed, 58 insertions(+), 8 deletions(-) diff --git a/src/agent/core/AgentCore.ts b/src/agent/core/AgentCore.ts index d404c2f..3dd9b7a 100644 --- a/src/agent/core/AgentCore.ts +++ b/src/agent/core/AgentCore.ts @@ -110,6 +110,7 @@ export class AgentCore { let assistantTextContent = ''; const accumulatedToolCalls: ToolCall[] = []; let finishReason = 'stop'; + let performanceInfo: StreamChunk['performanceInfo']; for await (const chunk of this.llmProvider.generateStream(conversationHistory, systemPrompt, tools)) { if (chunk.type === 'text') { @@ -120,6 +121,7 @@ export class AgentCore { accumulatedToolCalls.push(chunk.toolCall); } else if (chunk.type === 'done') { finishReason = chunk.finishReason ?? 'stop'; + performanceInfo = chunk.performanceInfo; } } @@ -163,10 +165,9 @@ export class AgentCore { // LLM finished with text response (stop reason) this.agentState.updateMessage(this.currentAssistantMessageId, assistantTextContent); continueLoop = false; + yield { type: 'done', content: '', finishReason, performanceInfo }; } } - - yield { type: 'done', content: '', finishReason: 'stop' }; } finally { this.currentUserMessageId = null; this.currentAssistantMessageId = null; diff --git a/src/agent/llm/LocalBrowserLLMProvider.ts b/src/agent/llm/LocalBrowserLLMProvider.ts index 9b8b9cb..fb39812 100644 --- a/src/agent/llm/LocalBrowserLLMProvider.ts +++ b/src/agent/llm/LocalBrowserLLMProvider.ts @@ -296,6 +296,14 @@ export class LocalBrowserLLMProvider implements LLMProvider { yield { type: 'tool_call', content: '', toolCall }; } - yield { type: 'done', content: '', finishReason }; + yield { + type: 'done', + content: '', + finishReason, + performanceInfo: { + prefillTps, + generationTps, + }, + }; } } diff --git a/src/agent/llm/StreamingTypes.ts b/src/agent/llm/StreamingTypes.ts index 738ca01..4600a5f 100644 --- a/src/agent/llm/StreamingTypes.ts +++ b/src/agent/llm/StreamingTypes.ts @@ -4,10 +4,16 @@ import type { ToolCall } from '../core/AgentState'; +export interface PerformanceInfo { + prefillTps?: number; + generationTps?: number; +} + export interface StreamChunk { type: 'text' | 'tool_call' | 'tool_result' | 'done'; content: string; toolCall?: ToolCall; toolResult?: { name: string; success: boolean; result: string }; - finishReason?: string; // 'stop' | 'tool_calls' — present on 'done' chunks + performanceInfo?: PerformanceInfo; + finishReason?: string; } diff --git a/src/components/ChatBox.css b/src/components/ChatBox.css index 4f61eec..6f63369 100644 --- a/src/components/ChatBox.css +++ b/src/components/ChatBox.css @@ -226,6 +226,12 @@ font-weight: bold; } +.message-performance-info { + margin-top: 8px; + font-size: 10px; + color: #909090; +} + /* Abort link styling */ .abort-link { background: none !important; @@ -310,4 +316,4 @@ 100% { background-position: -200% 0%; } -} \ No newline at end of file +} diff --git a/src/components/ChatBox.tsx b/src/components/ChatBox.tsx index 2b9940a..4a1cdfd 100644 --- a/src/components/ChatBox.tsx +++ b/src/components/ChatBox.tsx @@ -434,6 +434,7 @@ const ChatBox: React.FC = ({ isVisible }) => { key={message.id} content={message.content} isStreaming={message.isStreaming} + performanceInfo={message.performanceInfo} onAbort={message.isStreaming ? handleAbort : undefined} /> ) diff --git a/src/components/chat/AssistantMessage.tsx b/src/components/chat/AssistantMessage.tsx index 522ae4a..ad196be 100644 --- a/src/components/chat/AssistantMessage.tsx +++ b/src/components/chat/AssistantMessage.tsx @@ -3,11 +3,13 @@ import ReactMarkdown from 'react-markdown'; import remarkGfm from 'remark-gfm'; import { Prism as SyntaxHighlighter } from 'react-syntax-highlighter'; import { vscDarkPlus } from 'react-syntax-highlighter/dist/esm/styles/prism'; +import type { PerformanceInfo } from '../../agent/llm/StreamingTypes'; interface AssistantMessageProps { content: string; isStreaming?: boolean; onAbort?: () => void; + performanceInfo?: PerformanceInfo; } // Memoized code component to prevent SyntaxHighlighter re-renders @@ -30,7 +32,19 @@ const CodeComponent = memo(({ inline, className, children, ...props }: any) => { ); }); -const AssistantMessage: React.FC = ({ content, isStreaming, onAbort }) => { +const formatTps = (value?: number): string | null => { + if (typeof value !== 'number' || !Number.isFinite(value) || value <= 0) { + return null; + } + + return value.toFixed(1); +}; + +const AssistantMessage: React.FC = ({ content, isStreaming, onAbort, performanceInfo }) => { + const prefillTps = formatTps(performanceInfo?.prefillTps); + const generationTps = formatTps(performanceInfo?.generationTps); + const hasPerformanceInfo = Boolean(prefillTps || generationTps); + const renderContent = () => { // Handle special abort link for streaming messages if (isStreaming && onAbort && content.includes('click here to abort')) { @@ -83,6 +97,13 @@ const AssistantMessage: React.FC = ({ content, isStreamin
{renderContent()} + {hasPerformanceInfo && ( +
+ {prefillTps ? `Prefill: ${prefillTps} t/s` : 'Prefill: -'} + {' · '} + {generationTps ? `Generation: ${generationTps} t/s` : 'Generation: -'} +
+ )}
); diff --git a/src/hooks/useStreamProcessor.ts b/src/hooks/useStreamProcessor.ts index 31f6011..3a6c210 100644 --- a/src/hooks/useStreamProcessor.ts +++ b/src/hooks/useStreamProcessor.ts @@ -38,6 +38,7 @@ export const useStreamProcessor = (options: StreamProcessorOptions): StreamProce let assistantResponse = ''; let tokenCount = 0; let hasTextContent = false; + let performanceInfo: ChatMessage['performanceInfo']; console.log(`------------ ${logPrefix} ------------`); console.log(input); @@ -65,7 +66,8 @@ export const useStreamProcessor = (options: StreamProcessorOptions): StreamProce ...msg, content: assistantResponse, isStreaming: false, - tokenCount: undefined + tokenCount: undefined, + performanceInfo })); console.log('------------ ASSISTANT ------------'); @@ -98,19 +100,22 @@ export const useStreamProcessor = (options: StreamProcessorOptions): StreamProce assistantResponse = ''; tokenCount = 0; hasTextContent = false; + performanceInfo = undefined; // Create a fresh streaming placeholder for the next LLM response const nextMsg = createStreamingMessage(); currentStreamingId = nextMsg.id; onMessageAdd(nextMsg); } else if (chunk.type === 'done') { + performanceInfo = chunk.performanceInfo; // Finalize the streaming message if (hasTextContent) { onMessageUpdate(currentStreamingId, (msg) => ({ ...msg, content: assistantResponse, isStreaming: false, - tokenCount: undefined + tokenCount: undefined, + performanceInfo })); } else { // No text in final response — remove empty placeholder diff --git a/src/types/projectTypes.ts b/src/types/projectTypes.ts index 3c59d53..b94f5b4 100644 --- a/src/types/projectTypes.ts +++ b/src/types/projectTypes.ts @@ -1,4 +1,5 @@ import { Transform, type TransformFnParams } from 'class-transformer'; +import type { PerformanceInfo } from '../agent/llm/StreamingTypes'; export interface TimeSignature { numerator: number; @@ -11,6 +12,7 @@ export interface ChatMessage { content: string; isStreaming?: boolean; tokenCount?: number; + performanceInfo?: PerformanceInfo; } /**