chore: token speed and edit message (#5031)

* chore: add token speed measurement

* chore: add edit message handler

* chore: add DialogClose wrapper around save button

web-app/src/containers/ChatInput.tsx

@@ -49,6 +49,7 @@ const ChatInput = ({
   const { prompt, setPrompt } = usePrompt()
   const { t } = useTranslation()
   const { spellCheckChatInput } = useGeneralSetting()
+  const { tokenSpeed } = useAppState()
   const maxRows = 10
 
   const { selectedModel } = useModelProvider()
@@ -226,7 +227,7 @@ const ChatInput = ({
             {showSpeedToken && (
               <div className="flex items-center gap-1 text-main-view-fg/60 text-xs">
                 <IconBrandSpeedtest size={18} />
-                <span>42 tokens/sec</span>
+                <span>{Math.round(tokenSpeed?.tokenSpeed ?? 0)} tokens/sec</span>
               </div>
             )}
           </div>

web-app/src/containers/ThreadContent.tsx

@@ -14,6 +14,18 @@ import { useMessages } from '@/hooks/useMessages'
 import ThinkingBlock from '@/containers/ThinkingBlock'
 import ToolCallBlock from '@/containers/ToolCallBlock'
 import { useChat } from '@/hooks/useChat'
+import {
+  Dialog,
+  DialogClose,
+  DialogContent,
+  DialogFooter,
+  DialogHeader,
+  DialogTitle,
+  DialogTrigger,
+} from '@/components/ui/dialog'
+import { Button } from '@/components/ui/button'
+import { Textarea } from '@/components/ui/textarea'
+import { toast } from 'sonner'
 
 const CopyButton = ({ text }: { text: string }) => {
   const [copied, setCopied] = useState(false)
@@ -49,6 +61,8 @@ const CopyButton = ({ text }: { text: string }) => {
 // Use memo to prevent unnecessary re-renders, but allow re-renders when props change
 export const ThreadContent = memo(
   (item: ThreadMessage & { isLastMessage?: boolean; index?: number }) => {
+    const [message, setMessage] = useState(item.content?.[0]?.text?.value || '')
+
     // Use useMemo to stabilize the components prop
     const linkComponents = useMemo(
       () => ({
@@ -94,6 +108,20 @@ export const ThreadContent = memo(
       sendMessage(lastMessage.content?.[0]?.text?.value || '')
     }, [deleteMessage, getMessages, item, sendMessage])
 
+    const editMessage = useCallback(
+      (messageId: string) => {
+        const threadMessages = getMessages(item.thread_id)
+        const index = threadMessages.findIndex((msg) => msg.id === messageId)
+        if (index === -1) return
+        // Delete all messages after the edited message
+        for (let i = threadMessages.length - 1; i >= index; i--) {
+          deleteMessage(threadMessages[i].thread_id, threadMessages[i].id)
+        }
+        sendMessage(message)
+      },
+      [deleteMessage, getMessages, item.thread_id, message, sendMessage]
+    )
+
     const isToolCalls =
       item.metadata &&
       'tool_calls' in item.metadata &&
@@ -110,6 +138,8 @@ export const ThreadContent = memo(
               </div>
             </div>
             <div className="flex items-center justify-end gap-2 text-main-view-fg/60 text-xs mt-2">
+              <Dialog>
+                <DialogTrigger asChild>
                   <button
                     className="flex items-center gap-1 hover:text-accent transition-colors cursor-pointer group relative"
                     onClick={() => {
@@ -121,6 +151,50 @@ export const ThreadContent = memo(
                       Edit
                     </span>
                   </button>
+                </DialogTrigger>
+                <DialogContent>
+                  <DialogHeader>
+                    <DialogTitle>Edit Message</DialogTitle>
+                    <Textarea
+                      value={message}
+                      onChange={(e) => {
+                        setMessage(e.target.value)
+                      }}
+                      className="mt-2 resize-none"
+                      onKeyDown={(e) => {
+                        // Prevent key from being captured by parent components
+                        e.stopPropagation()
+                      }}
+                    />
+                    <DialogFooter className="mt-2 flex items-center">
+                      <DialogClose asChild>
+                        <Button
+                          variant="link"
+                          size="sm"
+                          className="hover:no-underline"
+                        >
+                          Cancel
+                        </Button>
+                      </DialogClose>
+                      <DialogClose asChild>
+                        <Button
+                          disabled={!message}
+                          onClick={() => {
+                            editMessage(item.id)
+                            toast.success('Edit Message', {
+                              id: 'edit-message',
+                              description:
+                                'Message edited successfully. Please wait for the model to respond.',
+                            })
+                          }}
+                        >
+                          Save
+                        </Button>
+                      </DialogClose>
+                    </DialogFooter>
+                  </DialogHeader>
+                </DialogContent>
+              </Dialog>
               <button
                 className="flex items-center gap-1 hover:text-accent transition-colors cursor-pointer group relative"
                 onClick={() => {

web-app/src/hooks/useAppState.ts

@@ -8,11 +8,14 @@ type AppState = {
   tools: MCPTool[]
   serverStatus: 'running' | 'stopped' | 'pending'
   abortControllers: Record<string, AbortController>
+  tokenSpeed?: TokenSpeed
   setServerStatus: (value: 'running' | 'stopped' | 'pending') => void
   updateStreamingContent: (content: ThreadMessage | undefined) => void
   updateLoadingModel: (loading: boolean) => void
   updateTools: (tools: MCPTool[]) => void
   setAbortController: (threadId: string, controller: AbortController) => void
+  updateTokenSpeed: (message: ThreadMessage) => void
+  resetTokenSpeed: () => void
 }
 
 export const useAppState = create<AppState>()((set) => ({
@@ -21,6 +24,7 @@ export const useAppState = create<AppState>()((set) => ({
   tools: [],
   serverStatus: 'stopped',
   abortControllers: {},
+  tokenSpeed: undefined,
   updateStreamingContent: (content) => {
     set({ streamingContent: content })
   },
@@ -39,4 +43,37 @@ export const useAppState = create<AppState>()((set) => ({
       },
     }))
   },
+  updateTokenSpeed: (message) =>
+    set((state) => {
+      const currentTimestamp = new Date().getTime() // Get current time in milliseconds
+      if (!state.tokenSpeed) {
+        // If this is the first update, just set the lastTimestamp and return
+        return {
+          tokenSpeed: {
+            lastTimestamp: currentTimestamp,
+            tokenSpeed: 0,
+            tokenCount: 1,
+            message: message.id,
+          },
+        }
+      }
+
+      const timeDiffInSeconds =
+        (currentTimestamp - state.tokenSpeed.lastTimestamp) / 1000 // Time difference in seconds
+      const totalTokenCount = state.tokenSpeed.tokenCount + 1
+      const averageTokenSpeed =
+        totalTokenCount / (timeDiffInSeconds > 0 ? timeDiffInSeconds : 1) // Calculate average token speed
+      return {
+        tokenSpeed: {
+          ...state.tokenSpeed,
+          tokenSpeed: averageTokenSpeed,
+          tokenCount: totalTokenCount,
+          message: message.id,
+        },
+      }
+    }),
+  resetTokenSpeed: () =>
+    set({
+      tokenSpeed: undefined,
+    }),
 }))

web-app/src/hooks/useChat.ts

@@ -22,7 +22,7 @@ import { useAssistant } from './useAssistant'
 
 export const useChat = () => {
   const { prompt, setPrompt } = usePrompt()
-  const { tools } = useAppState()
+  const { tools, updateTokenSpeed, resetTokenSpeed } = useAppState()
   const { currentAssistant } = useAssistant()
 
   const { getProviderByName, selectedModel, selectedProvider } =
@@ -68,6 +68,7 @@ export const useChat = () => {
     async (message: string) => {
       const activeThread = await getCurrentThread()
 
+      resetTokenSpeed()
       if (!activeThread || !provider) return
 
       updateStreamingContent(emptyThreadContent)
@@ -119,6 +120,7 @@ export const useChat = () => {
                 accumulatedText
               )
               updateStreamingContent(currentContent)
+              updateTokenSpeed(currentContent)
               await new Promise((resolve) => setTimeout(resolve, 0))
             }
           }
@@ -144,6 +146,7 @@ export const useChat = () => {
     },
     [
       getCurrentThread,
+      resetTokenSpeed,
       provider,
       updateStreamingContent,
       addMessage,
@@ -153,6 +156,7 @@ export const useChat = () => {
       setAbortController,
       updateLoadingModel,
       tools,
+      updateTokenSpeed,
     ]
   )
 

web-app/src/types/threads.d.ts

@@ -55,3 +55,10 @@ type Assistant = {
   instructions: string
   parameters: Record<string, unknown>
 }
+
+type TokenSpeed = {
+  message: string
+  tokenSpeed: number
+  tokenCount: number
+  lastTimestamp: number
+}