From 4cdab2b98b8625d1e013ada45ffc1f35b01b9fa8 Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Sun, 1 Dec 2024 16:11:37 +0700
Subject: [PATCH 1/5] fix: prevent other DOM elements from rerendering when
 generating

---
 .../SimpleTextMessage/index.tsx               | 446 +++++++++---------
 1 file changed, 230 insertions(+), 216 deletions(-)

diff --git a/web/screens/Thread/ThreadCenterPanel/SimpleTextMessage/index.tsx b/web/screens/Thread/ThreadCenterPanel/SimpleTextMessage/index.tsx
index d84ae57d7..a8d2cdbf2 100644
--- a/web/screens/Thread/ThreadCenterPanel/SimpleTextMessage/index.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/SimpleTextMessage/index.tsx
@@ -1,16 +1,11 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 /* eslint-disable react-hooks/exhaustive-deps */
 /* eslint-disable @typescript-eslint/naming-convention */
-import React, { useEffect, useState } from 'react'
+import React, { memo, useMemo } from 'react'
 
 import Markdown from 'react-markdown'
 
-import {
-  ChatCompletionRole,
-  ContentType,
-  MessageStatus,
-  ThreadMessage,
-} from '@janhq/core'
+import { ChatCompletionRole, ContentType, ThreadMessage } from '@janhq/core'
 
 import { Tooltip } from '@janhq/joi'
 
@@ -50,193 +45,29 @@ import {
 import { activeThreadAtom } from '@/helpers/atoms/Thread.atom'
 
 const SimpleTextMessage: React.FC<ThreadMessage> = (props) => {
-  let text = ''
   const isUser = props.role === ChatCompletionRole.User
   const isSystem = props.role === ChatCompletionRole.System
   const editMessage = useAtomValue(editMessageAtom)
   const activeThread = useAtomValue(activeThreadAtom)
 
-  if (props.content && props.content.length > 0) {
-    text = props.content[0]?.text?.value ?? ''
-  }
-
-  const clipboard = useClipboard({ timeout: 1000 })
-
-  function extractCodeLines(node: { children: { children: any[] }[] }) {
-    const codeLines: any[] = []
-
-    // Helper function to extract text recursively from children
-    function getTextFromNode(node: {
-      type: string
-      value: any
-      children: any[]
-    }): string {
-      if (node.type === 'text') {
-        return node.value
-      } else if (node.children) {
-        return node.children.map(getTextFromNode).join('')
-      }
-      return ''
-    }
-
-    // Traverse each line in the <code> block
-    node.children[0].children.forEach(
-      (lineNode: {
-        type: string
-        tagName: string
-        value: any
-        children: any[]
-      }) => {
-        if (lineNode.type === 'element' && lineNode.tagName === 'span') {
-          const lineContent = getTextFromNode(lineNode)
-          codeLines.push(lineContent)
-        }
-      }
-    )
-
-    // Join the lines with newline characters for proper formatting
-    return codeLines.join('\n')
-  }
-  function wrapCodeBlocksWithoutVisit() {
-    return (tree: { children: any[] }) => {
-      tree.children = tree.children.map((node) => {
-        if (node.tagName === 'pre' && node.children[0]?.tagName === 'code') {
-          const language = node.children[0].properties.className?.[1]?.replace(
-            'language-',
-            ''
-          )
-
-          if (extractCodeLines(node) === '') {
-            return node
-          }
-
-          return {
-            type: 'element',
-            tagName: 'div',
-            properties: {
-              className: ['code-block-wrapper'],
-            },
-            children: [
-              {
-                type: 'element',
-                tagName: 'div',
-                properties: {
-                  className: [
-                    'code-block',
-                    'group/item',
-                    'relative',
-                    'my-4',
-                    'overflow-auto',
-                  ],
-                },
-                children: [
-                  {
-                    type: 'element',
-                    tagName: 'div',
-                    properties: {
-                      className:
-                        'code-header bg-[hsla(var(--app-code-block))] flex justify-between items-center py-2 px-3 code-header--border rounded-t-lg',
-                    },
-                    children: [
-                      {
-                        type: 'element',
-                        tagName: 'span',
-                        properties: {
-                          className: 'text-xs font-medium text-gray-300',
-                        },
-                        children: [
-                          {
-                            type: 'text',
-                            value: language
-                              ? `${getLanguageFromExtension(language)}`
-                              : '',
-                          },
-                        ],
-                      },
-                      {
-                        type: 'element',
-                        tagName: 'button',
-                        properties: {
-                          className:
-                            'copy-button ml-auto flex items-center gap-1 text-xs font-medium text-gray-400 hover:text-gray-600 focus:outline-none',
-                          onClick: (event: Event) => {
-                            clipboard.copy(extractCodeLines(node))
-
-                            const button = event.currentTarget as HTMLElement
-                            button.innerHTML = `
-                              <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-check pointer-events-none text-green-600"><path d="M20 6 9 17l-5-5"/></svg>
-                              <span>Copied</span>
-                            `
-
-                            setTimeout(() => {
-                              button.innerHTML = `
-                                <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-copy pointer-events-none text-gray-400"><rect width="14" height="14" x="8" y="8" rx="2" ry="2"/><path d="M4 16c-1.1 0-2-.9-2-2V4c0-1.1.9-2 2-2h10c1.1 0 2 .9 2 2"/></svg>
-                                <span>Copy</span>
-                              `
-                            }, 2000)
-                          },
-                        },
-                        children: [
-                          {
-                            type: 'element',
-                            tagName: 'svg',
-                            properties: {
-                              xmlns: 'http://www.w3.org/2000/svg',
-                              width: '16',
-                              height: '16',
-                              viewBox: '0 0 24 24',
-                              fill: 'none',
-                              stroke: 'currentColor',
-                              strokeWidth: '2',
-                              strokeLinecap: 'round',
-                              strokeLinejoin: 'round',
-                              className:
-                                'lucide lucide-copy pointer-events-none text-gray-400',
-                            },
-                            children: [
-                              {
-                                type: 'element',
-                                tagName: 'rect',
-                                properties: {
-                                  width: '14',
-                                  height: '14',
-                                  x: '8',
-                                  y: '8',
-                                  rx: '2',
-                                  ry: '2',
-                                },
-                                children: [],
-                              },
-                              {
-                                type: 'element',
-                                tagName: 'path',
-                                properties: {
-                                  d: 'M4 16c-1.1 0-2-.9-2-2V4c0-1.1.9-2 2-2h10c1.1 0 2 .9 2 2',
-                                },
-                                children: [],
-                              },
-                            ],
-                          },
-                          { type: 'text', value: 'Copy' },
-                        ],
-                      },
-                    ],
-                  },
-                  node,
-                ],
-              },
-            ],
-          }
-        }
-        return node
-      })
-    }
-  }
-
   const { onViewFile, onViewFileContainer } = usePath()
   const tokenSpeed = useAtomValue(tokenSpeedAtom)
   const messages = useAtomValue(getCurrentChatMessagesAtom)
 
+  const text = useMemo(
+    () => props.content[0]?.text?.value ?? '',
+    [props.content]
+  )
+  const messageType = useMemo(
+    () => props.content[0]?.type ?? '',
+    [props.content]
+  )
+
+  const annotation = useMemo(
+    () => props.content[0]?.text?.annotations[0] ?? '',
+    [props.content]
+  )
+
   return (
     <div className="group relative mx-auto max-w-[700px] p-4">
       <div
@@ -302,15 +133,13 @@ const SimpleTextMessage: React.FC<ThreadMessage> = (props) => {
         )}
       >
         <>
-          {props.content[0]?.type === ContentType.Image && (
+          {messageType === ContentType.Image && (
             <div className="group/image relative mb-2 inline-flex cursor-pointer overflow-hidden rounded-xl">
               <div className="left-0 top-0 z-20 h-full w-full group-hover/image:inline-block">
                 <RelativeImage
-                  src={props.content[0]?.text.annotations[0]}
+                  src={annotation}
                   id={props.id}
-                  onClick={() =>
-                    onViewFile(`${props.content[0]?.text.annotations[0]}`)
-                  }
+                  onClick={() => onViewFile(annotation)}
                 />
               </div>
               <Tooltip
@@ -327,13 +156,11 @@ const SimpleTextMessage: React.FC<ThreadMessage> = (props) => {
             </div>
           )}
 
-          {props.content[0]?.type === ContentType.Pdf && (
+          {messageType === ContentType.Pdf && (
             <div className="group/file bg-secondary relative mb-2 inline-flex w-60 cursor-pointer gap-x-3 overflow-hidden rounded-lg p-4">
               <div
                 className="absolute left-0 top-0 z-20 hidden h-full w-full bg-black/20 backdrop-blur-sm group-hover/file:inline-block"
-                onClick={() =>
-                  onViewFile(`${props.id}.${props.content[0]?.type}`)
-                }
+                onClick={() => onViewFile(`${props.id}.${messageType}`)}
               />
               <Tooltip
                 trigger={
@@ -369,28 +196,10 @@ const SimpleTextMessage: React.FC<ThreadMessage> = (props) => {
               )}
               dir="ltr"
             >
-              <Markdown
-                remarkPlugins={[remarkMath]}
-                rehypePlugins={[
-                  [rehypeKatex, { throwOnError: false }],
-                  rehypeRaw,
-                  [
-                    // eslint-disable-next-line @typescript-eslint/ban-ts-comment
-                    // @ts-ignore
-                    rehypeHighlight,
-                    {
-                      languages: { latex },
-                      subset: false,
-                      plainText: ['txt', 'text'],
-                    },
-                  ],
-                  [rehypeHighlightCodeLines, { showLineNumbers: true }],
-                  wrapCodeBlocksWithoutVisit,
-                ]}
-                skipHtml={true}
-              >
-                {text}
-              </Markdown>
+              <MarkdownTextMessage
+                id={props.id}
+                text={text}
+              ></MarkdownTextMessage>
             </div>
           )}
         </>
@@ -399,4 +208,209 @@ const SimpleTextMessage: React.FC<ThreadMessage> = (props) => {
   )
 }
 
+export const MarkdownTextMessage = memo(
+  ({ text, id }: { id: string; text: string }) => {
+    console.log('rerender', id)
+    const clipboard = useClipboard({ timeout: 1000 })
+
+    function extractCodeLines(node: { children: { children: any[] }[] }) {
+      const codeLines: any[] = []
+
+      // Helper function to extract text recursively from children
+      function getTextFromNode(node: {
+        type: string
+        value: any
+        children: any[]
+      }): string {
+        if (node.type === 'text') {
+          return node.value
+        } else if (node.children) {
+          return node.children.map(getTextFromNode).join('')
+        }
+        return ''
+      }
+
+      // Traverse each line in the <code> block
+      node.children[0].children.forEach(
+        (lineNode: {
+          type: string
+          tagName: string
+          value: any
+          children: any[]
+        }) => {
+          if (lineNode.type === 'element' && lineNode.tagName === 'span') {
+            const lineContent = getTextFromNode(lineNode)
+            codeLines.push(lineContent)
+          }
+        }
+      )
+
+      // Join the lines with newline characters for proper formatting
+      return codeLines.join('\n')
+    }
+    function wrapCodeBlocksWithoutVisit() {
+      return (tree: { children: any[] }) => {
+        tree.children = tree.children.map((node) => {
+          if (node.tagName === 'pre' && node.children[0]?.tagName === 'code') {
+            const language =
+              node.children[0].properties.className?.[1]?.replace(
+                'language-',
+                ''
+              )
+
+            if (extractCodeLines(node) === '') {
+              return node
+            }
+
+            return {
+              type: 'element',
+              tagName: 'div',
+              properties: {
+                className: ['code-block-wrapper'],
+              },
+              children: [
+                {
+                  type: 'element',
+                  tagName: 'div',
+                  properties: {
+                    className: [
+                      'code-block',
+                      'group/item',
+                      'relative',
+                      'my-4',
+                      'overflow-auto',
+                    ],
+                  },
+                  children: [
+                    {
+                      type: 'element',
+                      tagName: 'div',
+                      properties: {
+                        className:
+                          'code-header bg-[hsla(var(--app-code-block))] flex justify-between items-center py-2 px-3 code-header--border rounded-t-lg',
+                      },
+                      children: [
+                        {
+                          type: 'element',
+                          tagName: 'span',
+                          properties: {
+                            className: 'text-xs font-medium text-gray-300',
+                          },
+                          children: [
+                            {
+                              type: 'text',
+                              value: language
+                                ? `${getLanguageFromExtension(language)}`
+                                : '',
+                            },
+                          ],
+                        },
+                        {
+                          type: 'element',
+                          tagName: 'button',
+                          properties: {
+                            className:
+                              'copy-button ml-auto flex items-center gap-1 text-xs font-medium text-gray-400 hover:text-gray-600 focus:outline-none',
+                            onClick: (event: Event) => {
+                              clipboard.copy(extractCodeLines(node))
+
+                              const button = event.currentTarget as HTMLElement
+                              button.innerHTML = `
+                              <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-check pointer-events-none text-green-600"><path d="M20 6 9 17l-5-5"/></svg>
+                              <span>Copied</span>
+                            `
+
+                              setTimeout(() => {
+                                button.innerHTML = `
+                                <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-copy pointer-events-none text-gray-400"><rect width="14" height="14" x="8" y="8" rx="2" ry="2"/><path d="M4 16c-1.1 0-2-.9-2-2V4c0-1.1.9-2 2-2h10c1.1 0 2 .9 2 2"/></svg>
+                                <span>Copy</span>
+                              `
+                              }, 2000)
+                            },
+                          },
+                          children: [
+                            {
+                              type: 'element',
+                              tagName: 'svg',
+                              properties: {
+                                xmlns: 'http://www.w3.org/2000/svg',
+                                width: '16',
+                                height: '16',
+                                viewBox: '0 0 24 24',
+                                fill: 'none',
+                                stroke: 'currentColor',
+                                strokeWidth: '2',
+                                strokeLinecap: 'round',
+                                strokeLinejoin: 'round',
+                                className:
+                                  'lucide lucide-copy pointer-events-none text-gray-400',
+                              },
+                              children: [
+                                {
+                                  type: 'element',
+                                  tagName: 'rect',
+                                  properties: {
+                                    width: '14',
+                                    height: '14',
+                                    x: '8',
+                                    y: '8',
+                                    rx: '2',
+                                    ry: '2',
+                                  },
+                                  children: [],
+                                },
+                                {
+                                  type: 'element',
+                                  tagName: 'path',
+                                  properties: {
+                                    d: 'M4 16c-1.1 0-2-.9-2-2V4c0-1.1.9-2 2-2h10c1.1 0 2 .9 2 2',
+                                  },
+                                  children: [],
+                                },
+                              ],
+                            },
+                            { type: 'text', value: 'Copy' },
+                          ],
+                        },
+                      ],
+                    },
+                    node,
+                  ],
+                },
+              ],
+            }
+          }
+          return node
+        })
+      }
+    }
+    return (
+      <>
+        <Markdown
+          remarkPlugins={[remarkMath]}
+          rehypePlugins={[
+            [rehypeKatex, { throwOnError: false }],
+            rehypeRaw,
+            [
+              // eslint-disable-next-line @typescript-eslint/ban-ts-comment
+              // @ts-ignore
+              rehypeHighlight,
+              {
+                languages: { latex },
+                subset: false,
+                plainText: ['txt', 'text'],
+              },
+            ],
+            [rehypeHighlightCodeLines, { showLineNumbers: true }],
+            wrapCodeBlocksWithoutVisit,
+          ]}
+          skipHtml={true}
+        >
+          {text}
+        </Markdown>
+      </>
+    )
+  }
+)
+
 export default React.memo(SimpleTextMessage)

From 0c8297d5c61b174b48ce6bb599e6744f4e7e194a Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Mon, 2 Dec 2024 09:48:48 +0700
Subject: [PATCH 2/5] chore: refactor message components

---
 .../ThreadCenterPanel/ChatItem/index.tsx      |  78 ++++++-
 .../MarkdownTextMessage.tsx}                  | 209 +-----------------
 .../RelativeImage.tsx                         |   0
 .../ThreadCenterPanel/TextMessage/index.tsx   | 129 +++++++++++
 4 files changed, 212 insertions(+), 204 deletions(-)
 rename web/screens/Thread/ThreadCenterPanel/{SimpleTextMessage/index.tsx => TextMessage/MarkdownTextMessage.tsx} (50%)
 rename web/screens/Thread/ThreadCenterPanel/{SimpleTextMessage => TextMessage}/RelativeImage.tsx (100%)
 create mode 100644 web/screens/Thread/ThreadCenterPanel/TextMessage/index.tsx

diff --git a/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx b/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx
index 92fd6394d..c7c05dca9 100644
--- a/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx
@@ -1,6 +1,7 @@
-import React, { forwardRef, useEffect, useState } from 'react'
+import React, { forwardRef, useEffect, useMemo, useState } from 'react'
 
 import {
+  ContentType,
   events,
   MessageEvent,
   MessageStatus,
@@ -8,9 +9,20 @@ import {
   ThreadMessage,
 } from '@janhq/core'
 
+import { Tooltip } from '@janhq/joi'
+
+import { FolderOpenIcon } from 'lucide-react'
+
 import ErrorMessage from '@/containers/ErrorMessage'
 
-import SimpleTextMessage from '../SimpleTextMessage'
+import { usePath } from '@/hooks/usePath'
+
+import { toGibibytes } from '@/utils/converter'
+import { openFileTitle } from '@/utils/titleUtils'
+
+import Icon from '../FileUploadPreview/Icon'
+import TextMessage from '../TextMessage'
+import { RelativeImage } from '../TextMessage/RelativeImage'
 
 type Ref = HTMLDivElement
 
@@ -27,6 +39,13 @@ const ChatItem = forwardRef<Ref, Props>((message, ref) => {
       ? message
       : undefined
   )
+  const messageType = useMemo(() => content[0]?.type ?? '', [content])
+
+  const annotation = useMemo(
+    () => content[0]?.text?.annotations[0] ?? '',
+    [content]
+  )
+  const { onViewFile, onViewFileContainer } = usePath()
 
   function onMessageUpdate(data: ThreadMessage) {
     if (data.id === message.id) {
@@ -54,7 +73,60 @@ const ChatItem = forwardRef<Ref, Props>((message, ref) => {
     <>
       {status !== MessageStatus.Error && content?.length > 0 && (
         <div ref={ref} className="relative">
-          <SimpleTextMessage {...message} content={content} status={status} />
+          {messageType === ContentType.Image && (
+            <div className="group/image relative mb-2 inline-flex cursor-pointer overflow-hidden rounded-xl">
+              <div className="left-0 top-0 z-20 h-full w-full group-hover/image:inline-block">
+                <RelativeImage
+                  src={annotation}
+                  id={message.id}
+                  onClick={() => onViewFile(annotation)}
+                />
+              </div>
+              <Tooltip
+                trigger={
+                  <div
+                    className="absolute right-2 top-2 z-20 hidden h-8 w-8 cursor-pointer items-center justify-center rounded-md bg-[hsla(var(--app-bg))] group-hover/image:flex"
+                    onClick={onViewFileContainer}
+                  >
+                    <FolderOpenIcon size={20} />
+                  </div>
+                }
+                content={<span>{openFileTitle()}</span>}
+              />
+            </div>
+          )}
+
+          {messageType === ContentType.Pdf && (
+            <div className="group/file bg-secondary relative mb-2 inline-flex w-60 cursor-pointer gap-x-3 overflow-hidden rounded-lg p-4">
+              <div
+                className="absolute left-0 top-0 z-20 hidden h-full w-full bg-black/20 backdrop-blur-sm group-hover/file:inline-block"
+                onClick={() => onViewFile(`${message.id}.${messageType}`)}
+              />
+              <Tooltip
+                trigger={
+                  <div
+                    className="absolute right-2 top-2 z-20 hidden h-8 w-8 cursor-pointer items-center justify-center rounded-md bg-[hsla(var(--app-bg))] group-hover/file:flex"
+                    onClick={onViewFileContainer}
+                  >
+                    <FolderOpenIcon size={20} />
+                  </div>
+                }
+                content={<span>{openFileTitle()}</span>}
+              />
+              <Icon type={content[0].type} />
+              <div className="w-full">
+                <h6 className="line-clamp-1 w-4/5 font-medium">
+                  {content[0].text.name?.replaceAll(/[-._]/g, ' ')}
+                </h6>
+                <p className="text-[hsla(var(--text-secondary)]">
+                  {toGibibytes(Number(content[0].text.size))}
+                </p>
+              </div>
+            </div>
+          )}
+          {messageType === ContentType.Text && (
+            <TextMessage {...message} content={content} status={status} />
+          )}
         </div>
       )}
       {errorMessage && !message.loadModelError && (
diff --git a/web/screens/Thread/ThreadCenterPanel/SimpleTextMessage/index.tsx b/web/screens/Thread/ThreadCenterPanel/TextMessage/MarkdownTextMessage.tsx
similarity index 50%
rename from web/screens/Thread/ThreadCenterPanel/SimpleTextMessage/index.tsx
rename to web/screens/Thread/ThreadCenterPanel/TextMessage/MarkdownTextMessage.tsx
index a8d2cdbf2..e9a6bba0f 100644
--- a/web/screens/Thread/ThreadCenterPanel/SimpleTextMessage/index.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/TextMessage/MarkdownTextMessage.tsx
@@ -1,217 +1,26 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 /* eslint-disable react-hooks/exhaustive-deps */
 /* eslint-disable @typescript-eslint/naming-convention */
-import React, { memo, useMemo } from 'react'
+import React, { memo } from 'react'
 
 import Markdown from 'react-markdown'
 
-import { ChatCompletionRole, ContentType, ThreadMessage } from '@janhq/core'
-
-import { Tooltip } from '@janhq/joi'
-
 import latex from 'highlight.js/lib/languages/latex'
-import { useAtomValue } from 'jotai'
-import { FolderOpenIcon } from 'lucide-react'
 import rehypeHighlight from 'rehype-highlight'
 import rehypeHighlightCodeLines from 'rehype-highlight-code-lines'
 import rehypeKatex from 'rehype-katex'
 import rehypeRaw from 'rehype-raw'
 import remarkMath from 'remark-math'
 import 'katex/dist/katex.min.css'
-import { twMerge } from 'tailwind-merge'
-
-import LogoMark from '@/containers/Brand/Logo/Mark'
 
 import { useClipboard } from '@/hooks/useClipboard'
-import { usePath } from '@/hooks/usePath'
 
 import { getLanguageFromExtension } from '@/utils/codeLanguageExtension'
-import { toGibibytes } from '@/utils/converter'
-import { displayDate } from '@/utils/datetime'
-
-import { openFileTitle } from '@/utils/titleUtils'
-
-import EditChatInput from '../EditChatInput'
-import Icon from '../FileUploadPreview/Icon'
-import MessageToolbar from '../MessageToolbar'
-
-import { RelativeImage } from './RelativeImage'
-
-import {
-  editMessageAtom,
-  getCurrentChatMessagesAtom,
-  tokenSpeedAtom,
-} from '@/helpers/atoms/ChatMessage.atom'
-import { activeThreadAtom } from '@/helpers/atoms/Thread.atom'
-
-const SimpleTextMessage: React.FC<ThreadMessage> = (props) => {
-  const isUser = props.role === ChatCompletionRole.User
-  const isSystem = props.role === ChatCompletionRole.System
-  const editMessage = useAtomValue(editMessageAtom)
-  const activeThread = useAtomValue(activeThreadAtom)
-
-  const { onViewFile, onViewFileContainer } = usePath()
-  const tokenSpeed = useAtomValue(tokenSpeedAtom)
-  const messages = useAtomValue(getCurrentChatMessagesAtom)
-
-  const text = useMemo(
-    () => props.content[0]?.text?.value ?? '',
-    [props.content]
-  )
-  const messageType = useMemo(
-    () => props.content[0]?.type ?? '',
-    [props.content]
-  )
-
-  const annotation = useMemo(
-    () => props.content[0]?.text?.annotations[0] ?? '',
-    [props.content]
-  )
-
-  return (
-    <div className="group relative mx-auto max-w-[700px] p-4">
-      <div
-        className={twMerge(
-          'mb-2 flex items-center justify-start gap-x-2',
-          !isUser && 'mt-2'
-        )}
-      >
-        {!isUser && !isSystem && <LogoMark width={28} />}
-        {isUser && (
-          <div className="flex h-8 w-8 items-center justify-center rounded-full border border-[hsla(var(--app-border))] last:border-none">
-            <svg
-              width="12"
-              height="16"
-              viewBox="0 0 12 16"
-              fill="none"
-              xmlns="http://www.w3.org/2000/svg"
-            >
-              <path
-                d="M6 0.497864C4.34315 0.497864 3 1.84101 3 3.49786C3 5.15472 4.34315 6.49786 6 6.49786C7.65685 6.49786 9 5.15472 9 3.49786C9 1.84101 7.65685 0.497864 6 0.497864ZM9.75 7.99786L2.24997 7.99787C1.00734 7.99787 0 9.00527 0 10.2479C0 11.922 0.688456 13.2633 1.81822 14.1701C2.93013 15.0625 4.42039 15.4979 6 15.4979C7.57961 15.4979 9.06987 15.0625 10.1818 14.1701C11.3115 13.2633 12 11.922 12 10.2479C12 9.00522 10.9926 7.99786 9.75 7.99786Z"
-                fill="#9CA3AF"
-              />
-            </svg>
-          </div>
-        )}
-
-        <div
-          className={twMerge(
-            'font-extrabold capitalize',
-            isUser && 'text-gray-500'
-          )}
-        >
-          {isUser
-            ? props.role
-            : (activeThread?.assistants[0].assistant_name ?? props.role)}
-        </div>
-        <p className="text-xs font-medium text-gray-400">
-          {displayDate(props.created)}
-        </p>
-        <div
-          className={twMerge(
-            'absolute right-0 cursor-pointer transition-all',
-            messages[messages.length - 1]?.id === props.id && !isUser
-              ? 'absolute -bottom-8 right-4'
-              : 'hidden group-hover:absolute group-hover:right-4 group-hover:top-4 group-hover:flex'
-          )}
-        >
-          <MessageToolbar message={props} />
-        </div>
-        {tokenSpeed &&
-          tokenSpeed.message === props.id &&
-          tokenSpeed.tokenSpeed > 0 && (
-            <p className="absolute right-8 text-xs font-medium text-[hsla(var(--text-secondary))]">
-              Token Speed: {Number(tokenSpeed.tokenSpeed).toFixed(2)}t/s
-            </p>
-          )}
-      </div>
-
-      <div
-        className={twMerge(
-          'w-full',
-          !isUser && !text.includes(' ') && 'break-all'
-        )}
-      >
-        <>
-          {messageType === ContentType.Image && (
-            <div className="group/image relative mb-2 inline-flex cursor-pointer overflow-hidden rounded-xl">
-              <div className="left-0 top-0 z-20 h-full w-full group-hover/image:inline-block">
-                <RelativeImage
-                  src={annotation}
-                  id={props.id}
-                  onClick={() => onViewFile(annotation)}
-                />
-              </div>
-              <Tooltip
-                trigger={
-                  <div
-                    className="absolute right-2 top-2 z-20 hidden h-8 w-8 cursor-pointer items-center justify-center rounded-md bg-[hsla(var(--app-bg))] group-hover/image:flex"
-                    onClick={onViewFileContainer}
-                  >
-                    <FolderOpenIcon size={20} />
-                  </div>
-                }
-                content={<span>{openFileTitle()}</span>}
-              />
-            </div>
-          )}
-
-          {messageType === ContentType.Pdf && (
-            <div className="group/file bg-secondary relative mb-2 inline-flex w-60 cursor-pointer gap-x-3 overflow-hidden rounded-lg p-4">
-              <div
-                className="absolute left-0 top-0 z-20 hidden h-full w-full bg-black/20 backdrop-blur-sm group-hover/file:inline-block"
-                onClick={() => onViewFile(`${props.id}.${messageType}`)}
-              />
-              <Tooltip
-                trigger={
-                  <div
-                    className="absolute right-2 top-2 z-20 hidden h-8 w-8 cursor-pointer items-center justify-center rounded-md bg-[hsla(var(--app-bg))] group-hover/file:flex"
-                    onClick={onViewFileContainer}
-                  >
-                    <FolderOpenIcon size={20} />
-                  </div>
-                }
-                content={<span>{openFileTitle()}</span>}
-              />
-              <Icon type={props.content[0].type} />
-              <div className="w-full">
-                <h6 className="line-clamp-1 w-4/5 font-medium">
-                  {props.content[0].text.name?.replaceAll(/[-._]/g, ' ')}
-                </h6>
-                <p className="text-[hsla(var(--text-secondary)]">
-                  {toGibibytes(Number(props.content[0].text.size))}
-                </p>
-              </div>
-            </div>
-          )}
-
-          {editMessage === props.id ? (
-            <div>
-              <EditChatInput message={props} />
-            </div>
-          ) : (
-            <div
-              className={twMerge(
-                'message max-width-[100%] flex flex-col gap-y-2 overflow-auto leading-relaxed'
-              )}
-              dir="ltr"
-            >
-              <MarkdownTextMessage
-                id={props.id}
-                text={text}
-              ></MarkdownTextMessage>
-            </div>
-          )}
-        </>
-      </div>
-    </div>
-  )
-}
 
 export const MarkdownTextMessage = memo(
   ({ text, id }: { id: string; text: string }) => {
-    console.log('rerender', id)
     const clipboard = useClipboard({ timeout: 1000 })
+    console.log('rerender', id)
 
     function extractCodeLines(node: { children: { children: any[] }[] }) {
       const codeLines: any[] = []
@@ -316,15 +125,15 @@ export const MarkdownTextMessage = memo(
 
                               const button = event.currentTarget as HTMLElement
                               button.innerHTML = `
-                              <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-check pointer-events-none text-green-600"><path d="M20 6 9 17l-5-5"/></svg>
-                              <span>Copied</span>
-                            `
+                                <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-check pointer-events-none text-green-600"><path d="M20 6 9 17l-5-5"/></svg>
+                                <span>Copied</span>
+                              `
 
                               setTimeout(() => {
                                 button.innerHTML = `
-                                <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-copy pointer-events-none text-gray-400"><rect width="14" height="14" x="8" y="8" rx="2" ry="2"/><path d="M4 16c-1.1 0-2-.9-2-2V4c0-1.1.9-2 2-2h10c1.1 0 2 .9 2 2"/></svg>
-                                <span>Copy</span>
-                              `
+                                  <svg xmlns="http://www.w3.org/2000/svg" width="16" height="16" viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round" class="lucide lucide-copy pointer-events-none text-gray-400"><rect width="14" height="14" x="8" y="8" rx="2" ry="2"/><path d="M4 16c-1.1 0-2-.9-2-2V4c0-1.1.9-2 2-2h10c1.1 0 2 .9 2 2"/></svg>
+                                  <span>Copy</span>
+                                `
                               }, 2000)
                             },
                           },
@@ -412,5 +221,3 @@ export const MarkdownTextMessage = memo(
     )
   }
 )
-
-export default React.memo(SimpleTextMessage)
diff --git a/web/screens/Thread/ThreadCenterPanel/SimpleTextMessage/RelativeImage.tsx b/web/screens/Thread/ThreadCenterPanel/TextMessage/RelativeImage.tsx
similarity index 100%
rename from web/screens/Thread/ThreadCenterPanel/SimpleTextMessage/RelativeImage.tsx
rename to web/screens/Thread/ThreadCenterPanel/TextMessage/RelativeImage.tsx
diff --git a/web/screens/Thread/ThreadCenterPanel/TextMessage/index.tsx b/web/screens/Thread/ThreadCenterPanel/TextMessage/index.tsx
new file mode 100644
index 000000000..d5e9030d3
--- /dev/null
+++ b/web/screens/Thread/ThreadCenterPanel/TextMessage/index.tsx
@@ -0,0 +1,129 @@
+import React, { useMemo } from 'react'
+
+import { ChatCompletionRole, ContentType, ThreadMessage } from '@janhq/core'
+
+import { Tooltip } from '@janhq/joi'
+
+import { useAtomValue } from 'jotai'
+import { FolderOpenIcon } from 'lucide-react'
+import 'katex/dist/katex.min.css'
+import { twMerge } from 'tailwind-merge'
+
+import LogoMark from '@/containers/Brand/Logo/Mark'
+
+import { usePath } from '@/hooks/usePath'
+
+import { displayDate } from '@/utils/datetime'
+
+import EditChatInput from '../EditChatInput'
+import MessageToolbar from '../MessageToolbar'
+
+import { MarkdownTextMessage } from './MarkdownTextMessage'
+
+import {
+  editMessageAtom,
+  getCurrentChatMessagesAtom,
+  tokenSpeedAtom,
+} from '@/helpers/atoms/ChatMessage.atom'
+import { activeThreadAtom } from '@/helpers/atoms/Thread.atom'
+
+const TextMessage: React.FC<ThreadMessage> = (props) => {
+  const isUser = props.role === ChatCompletionRole.User
+  const isSystem = props.role === ChatCompletionRole.System
+  const editMessage = useAtomValue(editMessageAtom)
+  const activeThread = useAtomValue(activeThreadAtom)
+
+  const tokenSpeed = useAtomValue(tokenSpeedAtom)
+  const messages = useAtomValue(getCurrentChatMessagesAtom)
+
+  const text = useMemo(
+    () => props.content[0]?.text?.value ?? '',
+    [props.content]
+  )
+
+  return (
+    <div className="group relative mx-auto max-w-[700px] p-4">
+      <div
+        className={twMerge(
+          'mb-2 flex items-center justify-start gap-x-2',
+          !isUser && 'mt-2'
+        )}
+      >
+        {!isUser && !isSystem && <LogoMark width={28} />}
+        {isUser && (
+          <div className="flex h-8 w-8 items-center justify-center rounded-full border border-[hsla(var(--app-border))] last:border-none">
+            <svg
+              width="12"
+              height="16"
+              viewBox="0 0 12 16"
+              fill="none"
+              xmlns="http://www.w3.org/2000/svg"
+            >
+              <path
+                d="M6 0.497864C4.34315 0.497864 3 1.84101 3 3.49786C3 5.15472 4.34315 6.49786 6 6.49786C7.65685 6.49786 9 5.15472 9 3.49786C9 1.84101 7.65685 0.497864 6 0.497864ZM9.75 7.99786L2.24997 7.99787C1.00734 7.99787 0 9.00527 0 10.2479C0 11.922 0.688456 13.2633 1.81822 14.1701C2.93013 15.0625 4.42039 15.4979 6 15.4979C7.57961 15.4979 9.06987 15.0625 10.1818 14.1701C11.3115 13.2633 12 11.922 12 10.2479C12 9.00522 10.9926 7.99786 9.75 7.99786Z"
+                fill="#9CA3AF"
+              />
+            </svg>
+          </div>
+        )}
+
+        <div
+          className={twMerge(
+            'font-extrabold capitalize',
+            isUser && 'text-gray-500'
+          )}
+        >
+          {isUser
+            ? props.role
+            : (activeThread?.assistants[0].assistant_name ?? props.role)}
+        </div>
+        <p className="text-xs font-medium text-gray-400">
+          {displayDate(props.created)}
+        </p>
+        <div
+          className={twMerge(
+            'absolute right-0 cursor-pointer transition-all',
+            messages[messages.length - 1]?.id === props.id && !isUser
+              ? 'absolute -bottom-8 right-4'
+              : 'hidden group-hover:absolute group-hover:right-4 group-hover:top-4 group-hover:flex'
+          )}
+        >
+          <MessageToolbar message={props} />
+        </div>
+        {tokenSpeed &&
+          tokenSpeed.message === props.id &&
+          tokenSpeed.tokenSpeed > 0 && (
+            <p className="absolute right-8 text-xs font-medium text-[hsla(var(--text-secondary))]">
+              Token Speed: {Number(tokenSpeed.tokenSpeed).toFixed(2)}t/s
+            </p>
+          )}
+      </div>
+
+      <div
+        className={twMerge(
+          'w-full',
+          !isUser && !text.includes(' ') && 'break-all'
+        )}
+      >
+        <>
+          {editMessage === props.id ? (
+            <div>
+              <EditChatInput message={props} />
+            </div>
+          ) : (
+            <div
+              className={twMerge(
+                'message max-width-[100%] flex flex-col gap-y-2 overflow-auto leading-relaxed'
+              )}
+              dir="ltr"
+            >
+              <MarkdownTextMessage id={props.id} text={text} />
+            </div>
+          )}
+        </>
+      </div>
+    </div>
+  )
+}
+
+export default React.memo(TextMessage)

From 19feb9e5061858aa36d6cb2acc3bc2b2ef0112b3 Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Mon, 2 Dec 2024 09:54:20 +0700
Subject: [PATCH 3/5] chore: clean imports

---
 web/screens/Thread/ThreadCenterPanel/TextMessage/index.tsx | 7 +------
 1 file changed, 1 insertion(+), 6 deletions(-)

diff --git a/web/screens/Thread/ThreadCenterPanel/TextMessage/index.tsx b/web/screens/Thread/ThreadCenterPanel/TextMessage/index.tsx
index d5e9030d3..905325e2e 100644
--- a/web/screens/Thread/ThreadCenterPanel/TextMessage/index.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/TextMessage/index.tsx
@@ -1,18 +1,13 @@
 import React, { useMemo } from 'react'
 
-import { ChatCompletionRole, ContentType, ThreadMessage } from '@janhq/core'
-
-import { Tooltip } from '@janhq/joi'
+import { ChatCompletionRole, ThreadMessage } from '@janhq/core'
 
 import { useAtomValue } from 'jotai'
-import { FolderOpenIcon } from 'lucide-react'
 import 'katex/dist/katex.min.css'
 import { twMerge } from 'tailwind-merge'
 
 import LogoMark from '@/containers/Brand/Logo/Mark'
 
-import { usePath } from '@/hooks/usePath'
-
 import { displayDate } from '@/utils/datetime'
 
 import EditChatInput from '../EditChatInput'

From e800349ed53d92789776f248dc05813e10ecbc97 Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Mon, 2 Dec 2024 09:57:29 +0700
Subject: [PATCH 4/5] chore: clean console log

---
 .../ThreadCenterPanel/TextMessage/MarkdownTextMessage.tsx      | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/web/screens/Thread/ThreadCenterPanel/TextMessage/MarkdownTextMessage.tsx b/web/screens/Thread/ThreadCenterPanel/TextMessage/MarkdownTextMessage.tsx
index e9a6bba0f..2b91271ae 100644
--- a/web/screens/Thread/ThreadCenterPanel/TextMessage/MarkdownTextMessage.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/TextMessage/MarkdownTextMessage.tsx
@@ -18,9 +18,8 @@ import { useClipboard } from '@/hooks/useClipboard'
 import { getLanguageFromExtension } from '@/utils/codeLanguageExtension'
 
 export const MarkdownTextMessage = memo(
-  ({ text, id }: { id: string; text: string }) => {
+  ({ text }: { id: string; text: string }) => {
     const clipboard = useClipboard({ timeout: 1000 })
-    console.log('rerender', id)
 
     function extractCodeLines(node: { children: { children: any[] }[] }) {
       const codeLines: any[] = []

From fc75fb64d255180a5ce83728343ccd4a65f7323e Mon Sep 17 00:00:00 2001
From: Louis <louis@jan.ai>
Date: Mon, 2 Dec 2024 10:35:51 +0700
Subject: [PATCH 5/5] chore: refactor message container and types

---
 .../ThreadCenterPanel/ChatItem/index.tsx      | 78 +------------------
 .../TextMessage/DocMessage.tsx                | 55 +++++++++++++
 .../TextMessage/ImageMessage.tsx              | 45 +++++++++++
 .../TextMessage/RelativeImage.tsx             |  4 +-
 .../ThreadCenterPanel/TextMessage/index.tsx   | 23 +++++-
 5 files changed, 124 insertions(+), 81 deletions(-)
 create mode 100644 web/screens/Thread/ThreadCenterPanel/TextMessage/DocMessage.tsx
 create mode 100644 web/screens/Thread/ThreadCenterPanel/TextMessage/ImageMessage.tsx

diff --git a/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx b/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx
index c7c05dca9..10d408211 100644
--- a/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/ChatItem/index.tsx
@@ -1,7 +1,6 @@
-import React, { forwardRef, useEffect, useMemo, useState } from 'react'
+import React, { forwardRef, useEffect, useState } from 'react'
 
 import {
-  ContentType,
   events,
   MessageEvent,
   MessageStatus,
@@ -9,20 +8,9 @@ import {
   ThreadMessage,
 } from '@janhq/core'
 
-import { Tooltip } from '@janhq/joi'
-
-import { FolderOpenIcon } from 'lucide-react'
-
 import ErrorMessage from '@/containers/ErrorMessage'
 
-import { usePath } from '@/hooks/usePath'
-
-import { toGibibytes } from '@/utils/converter'
-import { openFileTitle } from '@/utils/titleUtils'
-
-import Icon from '../FileUploadPreview/Icon'
-import TextMessage from '../TextMessage'
-import { RelativeImage } from '../TextMessage/RelativeImage'
+import MessageContainer from '../TextMessage'
 
 type Ref = HTMLDivElement
 
@@ -39,13 +27,6 @@ const ChatItem = forwardRef<Ref, Props>((message, ref) => {
       ? message
       : undefined
   )
-  const messageType = useMemo(() => content[0]?.type ?? '', [content])
-
-  const annotation = useMemo(
-    () => content[0]?.text?.annotations[0] ?? '',
-    [content]
-  )
-  const { onViewFile, onViewFileContainer } = usePath()
 
   function onMessageUpdate(data: ThreadMessage) {
     if (data.id === message.id) {
@@ -73,60 +54,7 @@ const ChatItem = forwardRef<Ref, Props>((message, ref) => {
     <>
       {status !== MessageStatus.Error && content?.length > 0 && (
         <div ref={ref} className="relative">
-          {messageType === ContentType.Image && (
-            <div className="group/image relative mb-2 inline-flex cursor-pointer overflow-hidden rounded-xl">
-              <div className="left-0 top-0 z-20 h-full w-full group-hover/image:inline-block">
-                <RelativeImage
-                  src={annotation}
-                  id={message.id}
-                  onClick={() => onViewFile(annotation)}
-                />
-              </div>
-              <Tooltip
-                trigger={
-                  <div
-                    className="absolute right-2 top-2 z-20 hidden h-8 w-8 cursor-pointer items-center justify-center rounded-md bg-[hsla(var(--app-bg))] group-hover/image:flex"
-                    onClick={onViewFileContainer}
-                  >
-                    <FolderOpenIcon size={20} />
-                  </div>
-                }
-                content={<span>{openFileTitle()}</span>}
-              />
-            </div>
-          )}
-
-          {messageType === ContentType.Pdf && (
-            <div className="group/file bg-secondary relative mb-2 inline-flex w-60 cursor-pointer gap-x-3 overflow-hidden rounded-lg p-4">
-              <div
-                className="absolute left-0 top-0 z-20 hidden h-full w-full bg-black/20 backdrop-blur-sm group-hover/file:inline-block"
-                onClick={() => onViewFile(`${message.id}.${messageType}`)}
-              />
-              <Tooltip
-                trigger={
-                  <div
-                    className="absolute right-2 top-2 z-20 hidden h-8 w-8 cursor-pointer items-center justify-center rounded-md bg-[hsla(var(--app-bg))] group-hover/file:flex"
-                    onClick={onViewFileContainer}
-                  >
-                    <FolderOpenIcon size={20} />
-                  </div>
-                }
-                content={<span>{openFileTitle()}</span>}
-              />
-              <Icon type={content[0].type} />
-              <div className="w-full">
-                <h6 className="line-clamp-1 w-4/5 font-medium">
-                  {content[0].text.name?.replaceAll(/[-._]/g, ' ')}
-                </h6>
-                <p className="text-[hsla(var(--text-secondary)]">
-                  {toGibibytes(Number(content[0].text.size))}
-                </p>
-              </div>
-            </div>
-          )}
-          {messageType === ContentType.Text && (
-            <TextMessage {...message} content={content} status={status} />
-          )}
+          <MessageContainer {...message} content={content} status={status} />
         </div>
       )}
       {errorMessage && !message.loadModelError && (
diff --git a/web/screens/Thread/ThreadCenterPanel/TextMessage/DocMessage.tsx b/web/screens/Thread/ThreadCenterPanel/TextMessage/DocMessage.tsx
new file mode 100644
index 000000000..9c0289734
--- /dev/null
+++ b/web/screens/Thread/ThreadCenterPanel/TextMessage/DocMessage.tsx
@@ -0,0 +1,55 @@
+import { memo } from 'react'
+
+import { Tooltip } from '@janhq/joi'
+
+import { FolderOpenIcon } from 'lucide-react'
+
+import { usePath } from '@/hooks/usePath'
+
+import { toGibibytes } from '@/utils/converter'
+import { openFileTitle } from '@/utils/titleUtils'
+
+import Icon from '../FileUploadPreview/Icon'
+
+const DocMessage = ({
+  id,
+  name,
+  size,
+}: {
+  id: string
+  name?: string
+  size?: number
+}) => {
+  const { onViewFile, onViewFileContainer } = usePath()
+
+  return (
+    <div className="group/file bg-secondary relative mb-2 inline-flex w-60 cursor-pointer gap-x-3 overflow-hidden rounded-lg p-4">
+      <div
+        className="absolute left-0 top-0 z-20 hidden h-full w-full bg-black/20 backdrop-blur-sm group-hover/file:inline-block"
+        onClick={() => onViewFile(`${id}.pdf`)}
+      />
+      <Tooltip
+        trigger={
+          <div
+            className="absolute right-2 top-2 z-20 hidden h-8 w-8 cursor-pointer items-center justify-center rounded-md bg-[hsla(var(--app-bg))] group-hover/file:flex"
+            onClick={onViewFileContainer}
+          >
+            <FolderOpenIcon size={20} />
+          </div>
+        }
+        content={<span>{openFileTitle()}</span>}
+      />
+      <Icon type="pdf" />
+      <div className="w-full">
+        <h6 className="line-clamp-1 w-4/5 font-medium">
+          {name?.replaceAll(/[-._]/g, ' ')}
+        </h6>
+        <p className="text-[hsla(var(--text-secondary)]">
+          {toGibibytes(Number(size))}
+        </p>
+      </div>
+    </div>
+  )
+}
+
+export default memo(DocMessage)
diff --git a/web/screens/Thread/ThreadCenterPanel/TextMessage/ImageMessage.tsx b/web/screens/Thread/ThreadCenterPanel/TextMessage/ImageMessage.tsx
new file mode 100644
index 000000000..117f259c0
--- /dev/null
+++ b/web/screens/Thread/ThreadCenterPanel/TextMessage/ImageMessage.tsx
@@ -0,0 +1,45 @@
+import { memo, useMemo } from 'react'
+
+import { ThreadContent } from '@janhq/core'
+import { Tooltip } from '@janhq/joi'
+
+import { FolderOpenIcon } from 'lucide-react'
+
+import { usePath } from '@/hooks/usePath'
+
+import { openFileTitle } from '@/utils/titleUtils'
+
+import { RelativeImage } from '../TextMessage/RelativeImage'
+
+const ImageMessage = ({ content }: { content: ThreadContent }) => {
+  const { onViewFile, onViewFileContainer } = usePath()
+
+  const annotation = useMemo(
+    () => content?.text?.annotations[0] ?? '',
+    [content]
+  )
+
+  return (
+    <div className="group/image relative mb-2 inline-flex cursor-pointer overflow-hidden rounded-xl">
+      <div className="left-0 top-0 z-20 h-full w-full group-hover/image:inline-block">
+        <RelativeImage
+          src={annotation}
+          onClick={() => onViewFile(annotation)}
+        />
+      </div>
+      <Tooltip
+        trigger={
+          <div
+            className="absolute right-2 top-2 z-20 hidden h-8 w-8 cursor-pointer items-center justify-center rounded-md bg-[hsla(var(--app-bg))] group-hover/image:flex"
+            onClick={onViewFileContainer}
+          >
+            <FolderOpenIcon size={20} />
+          </div>
+        }
+        content={<span>{openFileTitle()}</span>}
+      />
+    </div>
+  )
+}
+
+export default memo(ImageMessage)
diff --git a/web/screens/Thread/ThreadCenterPanel/TextMessage/RelativeImage.tsx b/web/screens/Thread/ThreadCenterPanel/TextMessage/RelativeImage.tsx
index fda4c0755..72d2a9365 100644
--- a/web/screens/Thread/ThreadCenterPanel/TextMessage/RelativeImage.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/TextMessage/RelativeImage.tsx
@@ -3,11 +3,9 @@ import { useEffect, useState } from 'react'
 import { getJanDataFolderPath } from '@janhq/core'
 
 export const RelativeImage = ({
-  id,
   src,
   onClick,
 }: {
-  id: string
   src: string
   onClick: () => void
 }) => {
@@ -22,7 +20,7 @@ export const RelativeImage = ({
     <button onClick={onClick}>
       <img
         className="aspect-auto h-[300px] cursor-pointer"
-        alt={id}
+        alt={src}
         src={src.includes('files/') ? `file://${path}/${src}` : src}
       />
     </button>
diff --git a/web/screens/Thread/ThreadCenterPanel/TextMessage/index.tsx b/web/screens/Thread/ThreadCenterPanel/TextMessage/index.tsx
index 905325e2e..7c4950b51 100644
--- a/web/screens/Thread/ThreadCenterPanel/TextMessage/index.tsx
+++ b/web/screens/Thread/ThreadCenterPanel/TextMessage/index.tsx
@@ -1,6 +1,6 @@
 import React, { useMemo } from 'react'
 
-import { ChatCompletionRole, ThreadMessage } from '@janhq/core'
+import { ChatCompletionRole, ContentType, ThreadMessage } from '@janhq/core'
 
 import { useAtomValue } from 'jotai'
 import 'katex/dist/katex.min.css'
@@ -13,6 +13,8 @@ import { displayDate } from '@/utils/datetime'
 import EditChatInput from '../EditChatInput'
 import MessageToolbar from '../MessageToolbar'
 
+import DocMessage from './DocMessage'
+import ImageMessage from './ImageMessage'
 import { MarkdownTextMessage } from './MarkdownTextMessage'
 
 import {
@@ -22,7 +24,7 @@ import {
 } from '@/helpers/atoms/ChatMessage.atom'
 import { activeThreadAtom } from '@/helpers/atoms/Thread.atom'
 
-const TextMessage: React.FC<ThreadMessage> = (props) => {
+const MessageContainer: React.FC<ThreadMessage> = (props) => {
   const isUser = props.role === ChatCompletionRole.User
   const isSystem = props.role === ChatCompletionRole.System
   const editMessage = useAtomValue(editMessageAtom)
@@ -35,6 +37,10 @@ const TextMessage: React.FC<ThreadMessage> = (props) => {
     () => props.content[0]?.text?.value ?? '',
     [props.content]
   )
+  const messageType = useMemo(
+    () => props.content[0]?.type ?? '',
+    [props.content]
+  )
 
   return (
     <div className="group relative mx-auto max-w-[700px] p-4">
@@ -101,6 +107,17 @@ const TextMessage: React.FC<ThreadMessage> = (props) => {
         )}
       >
         <>
+          {messageType === ContentType.Image && (
+            <ImageMessage content={props.content[0]} />
+          )}
+          {messageType === ContentType.Pdf && (
+            <DocMessage
+              id={props.id}
+              name={props.content[0]?.text?.name}
+              size={props.content[0]?.text?.size}
+            />
+          )}
+
           {editMessage === props.id ? (
             <div>
               <EditChatInput message={props} />
@@ -121,4 +138,4 @@ const TextMessage: React.FC<ThreadMessage> = (props) => {
   )
 }
 
-export default React.memo(TextMessage)
+export default React.memo(MessageContainer)