diff --git a/web/app/components/base/prompt-editor/plugins/component-picker-block/external-tool-option.tsx b/web/app/components/base/prompt-editor/plugins/component-picker-block/external-tool-option.tsx
deleted file mode 100644
index ffaf08a0f5..0000000000
--- a/web/app/components/base/prompt-editor/plugins/component-picker-block/external-tool-option.tsx
+++ /dev/null
@@ -1,85 +0,0 @@
-import { memo } from 'react'
-import { MenuOption } from '@lexical/react/LexicalTypeaheadMenuPlugin'
-
-export class VariableOption extends MenuOption {
-  title: string
-  icon?: JSX.Element
-  extraElement?: JSX.Element
-  keywords: Array<string>
-  keyboardShortcut?: string
-  onSelect: (queryString: string) => void
-
-  constructor(
-    title: string,
-    options: {
-      icon?: JSX.Element
-      extraElement?: JSX.Element
-      keywords?: Array<string>
-      keyboardShortcut?: string
-      onSelect: (queryString: string) => void
-    },
-  ) {
-    super(title)
-    this.title = title
-    this.keywords = options.keywords || []
-    this.icon = options.icon
-    this.extraElement = options.extraElement
-    this.keyboardShortcut = options.keyboardShortcut
-    this.onSelect = options.onSelect.bind(this)
-  }
-}
-
-type VariableMenuItemProps = {
-  isSelected: boolean
-  onClick: () => void
-  onMouseEnter: () => void
-  option: VariableOption
-  queryString: string | null
-}
-export const VariableMenuItem = memo(({
-  isSelected,
-  onClick,
-  onMouseEnter,
-  option,
-  queryString,
-}: VariableMenuItemProps) => {
-  const title = option.title
-  let before = title
-  let middle = ''
-  let after = ''
-
-  if (queryString) {
-    const regex = new RegExp(queryString, 'i')
-    const match = regex.exec(option.title)
-
-    if (match) {
-      before = title.substring(0, match.index)
-      middle = match[0]
-      after = title.substring(match.index + match[0].length)
-    }
-  }
-
-  return (
-    <div
-      key={option.key}
-      className={`
-        flex items-center px-3 h-6 rounded-md hover:bg-primary-50 cursor-pointer
-        ${isSelected && 'bg-primary-50'}
-      `}
-      tabIndex={-1}
-      ref={option.setRefElement}
-      onMouseEnter={onMouseEnter}
-      onClick={onClick}>
-      <div className='mr-2'>
-        {option.icon}
-      </div>
-      <div className='grow text-[13px] text-gray-900 truncate' title={option.title}>
-        {before}
-        <span className='text-[#2970FF]'>{middle}</span>
-        {after}
-      </div>
-      {option.extraElement}
-    </div>
-  )
-})
-VariableMenuItem.displayName = 'VariableMenuItem'
diff --git a/web/app/components/base/prompt-editor/plugins/component-picker-block/hooks.tsx b/web/app/components/base/prompt-editor/plugins/component-picker-block/hooks.tsx
index bb21a46366..b14bf8112b 100644
--- a/web/app/components/base/prompt-editor/plugins/component-picker-block/hooks.tsx
+++ b/web/app/components/base/prompt-editor/plugins/component-picker-block/hooks.tsx
@@ -15,8 +15,9 @@ import { INSERT_HISTORY_BLOCK_COMMAND } from '../history-block'
 import { INSERT_QUERY_BLOCK_COMMAND } from '../query-block'
 import { INSERT_VARIABLE_VALUE_BLOCK_COMMAND } from '../variable-block'
 import { $createCustomTextNode } from '../custom-text/node'
-import { PromptOption } from './prompt-option'
-import { VariableOption } from './variable-option'
+import { PromptMenuItem } from './prompt-option'
+import { VariableMenuItem } from './variable-option'
+import { PickerBlockMenuOption } from './menu'
 import { File05 } from '@/app/components/base/icons/src/vender/solid/files'
 import {
   MessageClockCircle,
@@ -35,62 +36,111 @@ export const usePromptOptions = (
   const { t } = useTranslation()
   const [editor] = useLexicalComposerContext()
 
-  return useMemo(() => {
-    return [
-      ...contextBlock?.show
-        ? [
-          new PromptOption(t('common.promptEditor.context.item.title'), {
-            icon: <File05 className='w-4 h-4 text-[#6938EF]' />,
-            onSelect: () => {
-              if (!contextBlock?.selectable)
-                return
-              editor.dispatchCommand(INSERT_CONTEXT_BLOCK_COMMAND, undefined)
-            },
-            disabled: !contextBlock?.selectable,
-          }),
-        ]
-        : [],
-      ...queryBlock?.show
-        ? [
-          new PromptOption(t('common.promptEditor.query.item.title'), {
-            icon: <UserEdit02 className='w-4 h-4 text-[#FD853A]' />,
-            onSelect: () => {
-              if (!queryBlock?.selectable)
-                return
-              editor.dispatchCommand(INSERT_QUERY_BLOCK_COMMAND, undefined)
-            },
-            disabled: !queryBlock?.selectable,
-          }),
-        ]
-        : [],
-      ...historyBlock?.show
-        ? [
-          new PromptOption(t('common.promptEditor.history.item.title'), {
-            icon: <MessageClockCircle className='w-4 h-4 text-[#DD2590]' />,
-            onSelect: () => {
-              if (!historyBlock?.selectable)
-                return
-              editor.dispatchCommand(INSERT_HISTORY_BLOCK_COMMAND, undefined)
-            },
-            disabled: !historyBlock?.selectable,
-          }),
-        ]
-        : [],
-    ]
-  }, [contextBlock, editor, historyBlock, queryBlock, t])
+  const promptOptions: PickerBlockMenuOption[] = []
+  if (contextBlock?.show) {
+    promptOptions.push(new PickerBlockMenuOption({
+      key: t('common.promptEditor.context.item.title'),
+      group: 'prompt context',
+      render: ({ isSelected, onSelect, onSetHighlight }) => {
+        return <PromptMenuItem
+          title={t('common.promptEditor.context.item.title')}
+          icon={<File05 className='w-4 h-4 text-[#6938EF]' />}
+          disabled={!contextBlock.selectable}
+          isSelected={isSelected}
+          onClick={onSelect}
+          onMouseEnter={onSetHighlight}
+        />
+      },
+      onSelect: () => {
+        if (!contextBlock?.selectable)
+          return
+        editor.dispatchCommand(INSERT_CONTEXT_BLOCK_COMMAND, undefined)
+      },
+    }))
+  }
+
+  if (queryBlock?.show) {
+    promptOptions.push(
+      new PickerBlockMenuOption({
+        key: t('common.promptEditor.query.item.title'),
+        group: 'prompt query',
+        render: ({ isSelected, onSelect, onSetHighlight }) => {
+          return (
+            <PromptMenuItem
+              title={t('common.promptEditor.query.item.title')}
+              icon={<UserEdit02 className='w-4 h-4 text-[#FD853A]' />}
+              disabled={!queryBlock.selectable}
+              isSelected={isSelected}
+              onClick={onSelect}
+              onMouseEnter={onSetHighlight}
+            />
+          )
+        },
+        onSelect: () => {
+          if (!queryBlock?.selectable)
+            return
+          editor.dispatchCommand(INSERT_QUERY_BLOCK_COMMAND, undefined)
+        },
+      }),
+    )
+  }
+
+  if (historyBlock?.show) {
+    promptOptions.push(
+      new PickerBlockMenuOption({
+        key: t('common.promptEditor.history.item.title'),
+        group: 'prompt history',
+        render: ({ isSelected, onSelect, onSetHighlight }) => {
+          return (
+            <PromptMenuItem
+              title={t('common.promptEditor.history.item.title')}
+              icon={<MessageClockCircle className='w-4 h-4 text-[#DD2590]' />}
+              disabled={!historyBlock.selectable
+              }
+              isSelected={isSelected}
+              onClick={onSelect}
+              onMouseEnter={onSetHighlight}
+            />
+          )
+        },
+        onSelect: () => {
+          if (!historyBlock?.selectable)
+            return
+          editor.dispatchCommand(INSERT_HISTORY_BLOCK_COMMAND, undefined)
+        },
+      }),
+    )
+  }
+  return promptOptions
 }
 
 export const useVariableOptions = (
   variableBlock?: VariableBlockType,
   queryString?: string,
-) => {
+): PickerBlockMenuOption[] => {
   const { t } = useTranslation()
   const [editor] = useLexicalComposerContext()
 
   const options = useMemo(() => {
-    const baseOptions = (variableBlock?.variables || []).map((item) => {
-      return new VariableOption(item.value, {
-        icon: <BracketsX className='w-[14px] h-[14px] text-[#2970FF]' />,
+    if (!variableBlock?.variables)
+      return []
+
+    const baseOptions = (variableBlock.variables).map((item) => {
+      return new PickerBlockMenuOption({
+        key: item.value,
+        group: 'prompt variable',
+        render: ({ queryString, isSelected, onSelect, onSetHighlight }) => {
+          return (
+            <VariableMenuItem
+              title={item.value}
+              icon={<BracketsX className='w-[14px] h-[14px] text-[#2970FF]' />}
+              queryString={queryString}
+              isSelected={isSelected}
+              onClick={onSelect}
+              onMouseEnter={onSetHighlight}
+            />
+          )
+        },
         onSelect: () => {
           editor.dispatchCommand(INSERT_VARIABLE_VALUE_BLOCK_COMMAND, `{{${item.value}}}`)
         },
@@ -101,12 +151,25 @@ export const useVariableOptions = (
 
     const regex = new RegExp(queryString, 'i')
 
-    return baseOptions.filter(option => regex.test(option.title) || option.keywords.some(keyword => regex.test(keyword)))
+    return baseOptions.filter(option => regex.test(option.key))
   }, [editor, queryString, variableBlock])
 
   const addOption = useMemo(() => {
-    return new VariableOption(t('common.promptEditor.variable.modal.add'), {
-      icon: <BracketsX className='mr-2 w-[14px] h-[14px] text-[#2970FF]' />,
+    return new PickerBlockMenuOption({
+      key: t('common.promptEditor.variable.modal.add'),
+      group: 'prompt variable',
+      render: ({ queryString, isSelected, onSelect, onSetHighlight }) => {
+        return (
+          <VariableMenuItem
+            title={t('common.promptEditor.variable.modal.add')}
+            icon={<BracketsX className='mr-2 w-[14px] h-[14px] text-[#2970FF]' />}
+            queryString={queryString}
+            isSelected={isSelected}
+            onClick={onSelect}
+            onMouseEnter={onSetHighlight}
+          />
+        )
+      },
       onSelect: () => {
         editor.update(() => {
           const prefixNode = $createCustomTextNode('{{')
@@ -131,16 +194,31 @@ export const useExternalToolOptions = (
   const [editor] = useLexicalComposerContext()
 
   const options = useMemo(() => {
-    const baseToolOptions = (externalToolBlockType?.externalTools || []).map((item) => {
-      return new VariableOption(item.name, {
-        icon: (
-          <AppIcon
-            className='!w-[14px] !h-[14px]'
-            icon={item.icon}
-            background={item.icon_background}
-          />
-        ),
-        extraElement: <div className='text-xs text-gray-400'>{item.variableName}</div>,
+    if (!externalToolBlockType?.externalTools)
+      return []
+    const baseToolOptions = (externalToolBlockType.externalTools).map((item) => {
+      return new PickerBlockMenuOption({
+        key: item.name,
+        group: 'external tool',
+        render: ({ queryString, isSelected, onSelect, onSetHighlight }) => {
+          return (
+            <VariableMenuItem
+              title={item.name}
+              icon={
+                <AppIcon
+                  className='!w-[14px] !h-[14px]'
+                  icon={item.icon}
+                  background={item.icon_background}
+                />
+              }
+              extraElement={<div className='text-xs text-gray-400'>{item.variableName}</div>}
+              queryString={queryString}
+              isSelected={isSelected}
+              onClick={onSelect}
+              onMouseEnter={onSetHighlight}
+            />
+          )
+        },
         onSelect: () => {
           editor.dispatchCommand(INSERT_VARIABLE_VALUE_BLOCK_COMMAND, `{{${item.variableName}}}`)
         },
@@ -151,16 +229,28 @@ export const useExternalToolOptions = (
 
     const regex = new RegExp(queryString, 'i')
 
-    return baseToolOptions.filter(option => regex.test(option.title) || option.keywords.some(keyword => regex.test(keyword)))
+    return baseToolOptions.filter(option => regex.test(option.key))
   }, [editor, queryString, externalToolBlockType])
 
   const addOption = useMemo(() => {
-    return new VariableOption(t('common.promptEditor.variable.modal.addTool'), {
-      icon: <Tool03 className='mr-2 w-[14px] h-[14px] text-[#444CE7]' />,
-      extraElement: <ArrowUpRight className='w-3 h-3 text-gray-400' />,
+    return new PickerBlockMenuOption({
+      key: t('common.promptEditor.variable.modal.addTool'),
+      group: 'external tool',
+      render: ({ queryString, isSelected, onSelect, onSetHighlight }) => {
+        return (
+          <VariableMenuItem
+            title={t('common.promptEditor.variable.modal.addTool')}
+            icon={<Tool03 className='mr-2 w-[14px] h-[14px] text-[#444CE7]' />}
+            extraElement={< ArrowUpRight className='w-3 h-3 text-gray-400' />}
+            queryString={queryString}
+            isSelected={isSelected}
+            onClick={onSelect}
+            onMouseEnter={onSetHighlight}
+          />
+        )
+      },
       onSelect: () => {
-        if (externalToolBlockType?.onAddExternalTool)
-          externalToolBlockType.onAddExternalTool()
+        externalToolBlockType?.onAddExternalTool?.()
       },
     })
   }, [externalToolBlockType, t])
@@ -191,11 +281,8 @@ export const useOptions = (
 
   return useMemo(() => {
     return {
-      promptOptions,
-      variableOptions,
-      externalToolOptions,
       workflowVariableOptions,
-      allOptions: [...promptOptions, ...variableOptions, ...externalToolOptions],
+      allFlattenOptions: [...promptOptions, ...variableOptions, ...externalToolOptions],
     }
   }, [promptOptions, variableOptions, externalToolOptions, workflowVariableOptions])
 }
diff --git a/web/app/components/base/prompt-editor/plugins/component-picker-block/index.tsx b/web/app/components/base/prompt-editor/plugins/component-picker-block/index.tsx
index 2500f72e8b..15b07ded17 100644
--- a/web/app/components/base/prompt-editor/plugins/component-picker-block/index.tsx
+++ b/web/app/components/base/prompt-editor/plugins/component-picker-block/index.tsx
@@ -1,11 +1,11 @@
 import {
+  Fragment,
   memo,
   useCallback,
   useState,
 } from 'react'
 import ReactDOM from 'react-dom'
 import {
-  FloatingPortal,
   flip,
   offset,
   shift,
@@ -27,11 +27,8 @@ import { useBasicTypeaheadTriggerMatch } from '../../hooks'
 import { INSERT_WORKFLOW_VARIABLE_BLOCK_COMMAND } from '../workflow-variable-block'
 import { INSERT_VARIABLE_VALUE_BLOCK_COMMAND } from '../variable-block'
 import { $splitNodeContainingQuery } from '../../utils'
-import type { PromptOption } from './prompt-option'
-import PromptMenu from './prompt-menu'
-import VariableMenu from './variable-menu'
-import type { VariableOption } from './variable-option'
 import { useOptions } from './hooks'
+import type { PickerBlockMenuOption } from './menu'
 import VarReferenceVars from '@/app/components/workflow/nodes/_base/components/variable/var-reference-vars'
 import { useEventEmitterContextContext } from '@/context/event-emitter'
 
@@ -54,11 +51,13 @@ const ComponentPicker = ({
   workflowVariableBlock,
 }: ComponentPickerProps) => {
   const { eventEmitter } = useEventEmitterContextContext()
-  const { refs, floatingStyles, elements } = useFloating({
+  const { refs, floatingStyles, isPositioned } = useFloating({
     placement: 'bottom-start',
     middleware: [
       offset(0), // fix hide cursor
-      shift(),
+      shift({
+        padding: 8,
+      }),
       flip(),
     ],
   })
@@ -76,10 +75,7 @@ const ComponentPicker = ({
   })
 
   const {
-    allOptions,
-    promptOptions,
-    variableOptions,
-    externalToolOptions,
+    allFlattenOptions,
     workflowVariableOptions,
   } = useOptions(
     contextBlock,
@@ -92,18 +88,15 @@ const ComponentPicker = ({
 
   const onSelectOption = useCallback(
     (
-      selectedOption: PromptOption | VariableOption,
+      selectedOption: PickerBlockMenuOption,
       nodeToRemove: TextNode | null,
       closeMenu: () => void,
-      matchingString: string,
     ) => {
       editor.update(() => {
         if (nodeToRemove && selectedOption?.key)
           nodeToRemove.remove()
 
-        if (selectedOption?.onSelect)
-          selectedOption.onSelect(matchingString)
-
+        selectedOption.onSelectMenuOption()
         closeMenu()
       })
     },
@@ -123,157 +116,93 @@ const ComponentPicker = ({
       editor.dispatchCommand(INSERT_WORKFLOW_VARIABLE_BLOCK_COMMAND, variables)
   }, [editor, checkForTriggerMatch, triggerString])
 
-  const renderMenu = useCallback<MenuRenderFn<PromptOption | VariableOption>>((
+  const renderMenu = useCallback<MenuRenderFn<PickerBlockMenuOption>>((
     anchorElementRef,
-    { selectedIndex, selectOptionAndCleanUp, setHighlightedIndex },
+    { options, selectedIndex, selectOptionAndCleanUp, setHighlightedIndex },
   ) => {
-    if (anchorElementRef.current && (allOptions.length || workflowVariableBlock?.show)) {
-      return (
-        <>
-          {
-            ReactDOM.createPortal(
-              <div ref={refs.setReference}></div>,
-              anchorElementRef.current,
-            )
-          }
-          {
-            elements.reference && (
-              <FloatingPortal id='typeahead-menu'>
-                <div
-                  className='w-[260px] bg-white rounded-lg border-[0.5px] border-gray-200 shadow-lg overflow-y-auto'
-                  style={{
-                    ...floatingStyles,
-                    maxHeight: 'calc(1 / 3 * 100vh)',
-                  }}
-                  ref={refs.setFloating}
-                >
-                  {
-                    !!promptOptions.length && (
-                      <>
-                        <PromptMenu
-                          startIndex={0}
-                          selectedIndex={selectedIndex}
-                          options={promptOptions}
-                          onClick={(index, option) => {
-                            if (option.disabled)
-                              return
-                            setHighlightedIndex(index)
-                            selectOptionAndCleanUp(option)
-                          }}
-                          onMouseEnter={(index, option) => {
-                            if (option.disabled)
-                              return
-                            setHighlightedIndex(index)
-                          }}
-                        />
-                      </>
-                    )
-                  }
-                  {
-                    !!variableOptions.length && (
-                      <>
-                        {
-                          !!promptOptions.length && (
-                            <div className='h-[1px] bg-gray-100'></div>
-                          )
-                        }
-                        <VariableMenu
-                          startIndex={promptOptions.length}
-                          selectedIndex={selectedIndex}
-                          options={variableOptions}
-                          onClick={(index, option) => {
-                            if (option.disabled)
-                              return
-                            setHighlightedIndex(index)
-                            selectOptionAndCleanUp(option)
-                          }}
-                          onMouseEnter={(index, option) => {
-                            if (option.disabled)
-                              return
-                            setHighlightedIndex(index)
-                          }}
-                          queryString={queryString}
-                        />
-                      </>
-                    )
-                  }
-                  {
-                    !!externalToolOptions.length && (
-                      <>
-                        {
-                          (!!promptOptions.length || !!variableOptions.length) && (
-                            <div className='h-[1px] bg-gray-100'></div>
-                          )
-                        }
-                        <VariableMenu
-                          startIndex={promptOptions.length + variableOptions.length}
-                          selectedIndex={selectedIndex}
-                          options={externalToolOptions}
-                          onClick={(index, option) => {
-                            if (option.disabled)
-                              return
-                            setHighlightedIndex(index)
-                            selectOptionAndCleanUp(option)
-                          }}
-                          onMouseEnter={(index, option) => {
-                            if (option.disabled)
-                              return
-                            setHighlightedIndex(index)
-                          }}
-                          queryString={queryString}
-                        />
-                      </>
-                    )
-                  }
-                  {
-                    workflowVariableBlock?.show && (
-                      <>
-                        {
-                          (!!promptOptions.length || !!variableOptions.length || !!externalToolOptions.length) && (
-                            <div className='h-[1px] bg-gray-100'></div>
-                          )
-                        }
-                        <div className='p-1'>
-                          <VarReferenceVars
-                            hideSearch
-                            vars={workflowVariableOptions}
-                            onChange={(variables: string[]) => {
-                              handleSelectWorkflowVariable(variables)
-                            }}
-                          />
-                        </div>
-                      </>
-                    )
-                  }
-                </div>
-              </FloatingPortal>
-            )
-          }
-        </>
-      )
-    }
+    if (!(anchorElementRef.current && (allFlattenOptions.length || workflowVariableBlock?.show)))
+      return null
+    refs.setReference(anchorElementRef.current)
 
-    return null
-  }, [
-    allOptions,
-    promptOptions,
-    variableOptions,
-    externalToolOptions,
-    queryString,
-    workflowVariableBlock?.show,
-    workflowVariableOptions,
-    handleSelectWorkflowVariable,
-    elements,
-    floatingStyles,
-    refs,
-  ])
+    return (
+      <>
+        {
+          ReactDOM.createPortal(
+            // The `LexicalMenu` will try to calculate the position of the floating menu based on the first child.
+            // Since we use floating ui, we need to wrap it with a div to prevent the position calculation being affected.
+            // See https://github.com/facebook/lexical/blob/ac97dfa9e14a73ea2d6934ff566282d7f758e8bb/packages/lexical-react/src/shared/LexicalMenu.ts#L493
+            <div className='w-0 h-0'>
+              <div
+                className='p-1 w-[260px] bg-white rounded-lg border-[0.5px] border-gray-200 shadow-lg overflow-y-auto overflow-x-hidden'
+                style={{
+                  ...floatingStyles,
+                  visibility: isPositioned ? 'visible' : 'hidden',
+                  maxHeight: 'calc(1 / 3 * 100vh)',
+                }}
+                ref={refs.setFloating}
+              >
+                {
+                  options.map((option, index) => (
+                    <Fragment key={option.key}>
+                      {
+                        // Divider
+                        index !== 0 && options.at(index - 1)?.group !== option.group && (
+                          <div className='h-px bg-gray-100 my-1 w-screen -translate-x-1'></div>
+                        )
+                      }
+                      {option.renderMenuOption({
+                        queryString,
+                        isSelected: selectedIndex === index,
+                        onSelect: () => {
+                          selectOptionAndCleanUp(option)
+                        },
+                        onSetHighlight: () => {
+                          setHighlightedIndex(index)
+                        },
+                      })}
+                    </Fragment>
+                  ))
+                }
+                {
+                  workflowVariableBlock?.show && (
+                    <>
+                      {
+                        (!!options.length) && (
+                          <div className='h-px bg-gray-100 my-1 w-screen -translate-x-1'></div>
+                        )
+                      }
+                      <div className='p-1'>
+                        <VarReferenceVars
+                          hideSearch
+                          vars={workflowVariableOptions}
+                          onChange={(variables: string[]) => {
+                            handleSelectWorkflowVariable(variables)
+                          }}
+                        />
+                      </div>
+                    </>
+                  )
+                }
+              </div>
+            </div>,
+            anchorElementRef.current,
+          )
+        }
+      </>
+    )
+  }, [allFlattenOptions.length, workflowVariableBlock?.show, refs, isPositioned, floatingStyles, queryString, workflowVariableOptions, handleSelectWorkflowVariable])
 
   return (
     <LexicalTypeaheadMenuPlugin
-      options={allOptions as any}
+      options={allFlattenOptions}
       onQueryChange={setQueryString}
       onSelectOption={onSelectOption}
-      anchorClassName='z-[999999]'
+      // The `translate` class is used to workaround the issue that the `typeahead-menu` menu is not positioned as expected.
+      // See also https://github.com/facebook/lexical/blob/772520509308e8ba7e4a82b6cd1996a78b3298d0/packages/lexical-react/src/shared/LexicalMenu.ts#L498
+      //
+      // We no need the position function of the `LexicalTypeaheadMenuPlugin`,
+      // so the reference anchor should be positioned based on the range of the trigger string, and the menu will be positioned by the floating ui.
+      anchorClassName='z-[999999] translate-y-[calc(-100%-3px)]'
       menuRenderFn={renderMenu}
       triggerFn={checkForTriggerMatch}
     />
diff --git a/web/app/components/base/prompt-editor/plugins/component-picker-block/menu.tsx b/web/app/components/base/prompt-editor/plugins/component-picker-block/menu.tsx
new file mode 100644
index 0000000000..d8c7156926
--- /dev/null
+++ b/web/app/components/base/prompt-editor/plugins/component-picker-block/menu.tsx
@@ -0,0 +1,31 @@
+import { MenuOption } from '@lexical/react/LexicalTypeaheadMenuPlugin'
+import { Fragment } from 'react'
+
+/**
+ * Corresponds to the `MenuRenderFn` type from `@lexical/react/LexicalTypeaheadMenuPlugin`.
+ */
+type MenuOptionRenderProps = {
+  isSelected: boolean
+  onSelect: () => void
+  onSetHighlight: () => void
+  queryString: string | null
+}
+
+export class PickerBlockMenuOption extends MenuOption {
+  public group?: string
+
+  constructor(
+    private data: {
+      key: string
+      group?: string
+      onSelect?: () => void
+      render: (menuRenderProps: MenuOptionRenderProps) => JSX.Element
+    },
+  ) {
+    super(data.key)
+    this.group = data.group
+  }
+
+  public onSelectMenuOption = () => this.data.onSelect?.()
+  public renderMenuOption = (menuRenderProps: MenuOptionRenderProps) => <Fragment key={this.data.key}>{this.data.render(menuRenderProps)}</Fragment>
+}
diff --git a/web/app/components/base/prompt-editor/plugins/component-picker-block/prompt-menu.tsx b/web/app/components/base/prompt-editor/plugins/component-picker-block/prompt-menu.tsx
deleted file mode 100644
index 6f16fcc2ba..0000000000
--- a/web/app/components/base/prompt-editor/plugins/component-picker-block/prompt-menu.tsx
+++ /dev/null
@@ -1,37 +0,0 @@
-import { memo } from 'react'
-import { PromptMenuItem } from './prompt-option'
-
-type PromptMenuProps = {
-  startIndex: number
-  selectedIndex: number | null
-  options: any[]
-  onClick: (index: number, option: any) => void
-  onMouseEnter: (index: number, option: any) => void
-}
-const PromptMenu = ({
-  startIndex,
-  selectedIndex,
-  options,
-  onClick,
-  onMouseEnter,
-}: PromptMenuProps) => {
-  return (
-    <div className='p-1'>
-      {
-        options.map((option, index: number) => (
-          <PromptMenuItem
-            startIndex={startIndex}
-            index={index}
-            isSelected={selectedIndex === index + startIndex}
-            onClick={onClick}
-            onMouseEnter={onMouseEnter}
-            key={option.key}
-            option={option}
-          />
-        ))
-      }
-    </div>
-  )
-}
-
-export default memo(PromptMenu)
diff --git a/web/app/components/base/prompt-editor/plugins/component-picker-block/prompt-option.tsx b/web/app/components/base/prompt-editor/plugins/component-picker-block/prompt-option.tsx
index 6937872786..7aabbe4b26 100644
--- a/web/app/components/base/prompt-editor/plugins/component-picker-block/prompt-option.tsx
+++ b/web/app/components/base/prompt-editor/plugins/component-picker-block/prompt-option.tsx
@@ -1,64 +1,44 @@
 import { memo } from 'react'
-import { MenuOption } from '@lexical/react/LexicalTypeaheadMenuPlugin'
-
-export class PromptOption extends MenuOption {
-  title: string
-  icon?: JSX.Element
-  keywords: Array<string>
-  keyboardShortcut?: string
-  onSelect: (queryString: string) => void
-  disabled?: boolean
-
-  constructor(
-    title: string,
-    options: {
-      icon?: JSX.Element
-      keywords?: Array<string>
-      keyboardShortcut?: string
-      onSelect: (queryString: string) => void
-      disabled?: boolean
-    },
-  ) {
-    super(title)
-    this.title = title
-    this.keywords = options.keywords || []
-    this.icon = options.icon
-    this.keyboardShortcut = options.keyboardShortcut
-    this.onSelect = options.onSelect.bind(this)
-    this.disabled = options.disabled
-  }
-}
 
 type PromptMenuItemMenuItemProps = {
-  startIndex: number
-  index: number
+  icon: JSX.Element
+  title: string
+  disabled?: boolean
   isSelected: boolean
-  onClick: (index: number, option: PromptOption) => void
-  onMouseEnter: (index: number, option: PromptOption) => void
-  option: PromptOption
+  onClick: () => void
+  onMouseEnter: () => void
+  setRefElement?: (element: HTMLDivElement) => void
 }
 export const PromptMenuItem = memo(({
-  startIndex,
-  index,
+  icon,
+  title,
+  disabled,
   isSelected,
   onClick,
   onMouseEnter,
-  option,
+  setRefElement,
 }: PromptMenuItemMenuItemProps) => {
   return (
     <div
-      key={option.key}
       className={`
         flex items-center px-3 h-6 cursor-pointer hover:bg-gray-50 rounded-md
-        ${isSelected && !option.disabled && '!bg-gray-50'}
-        ${option.disabled ? 'cursor-not-allowed opacity-30' : 'hover:bg-gray-50 cursor-pointer'}
+        ${isSelected && !disabled && '!bg-gray-50'}
+        ${disabled ? 'cursor-not-allowed opacity-30' : 'hover:bg-gray-50 cursor-pointer'}
       `}
       tabIndex={-1}
-      ref={option.setRefElement}
-      onMouseEnter={() => onMouseEnter(index + startIndex, option)}
-      onClick={() => onClick(index + startIndex, option)}>
-      {option.icon}
-      <div className='ml-1 text-[13px] text-gray-900'>{option.title}</div>
+      ref={setRefElement}
+      onMouseEnter={() => {
+        if (disabled)
+          return
+        onMouseEnter()
+      }}
+      onClick={() => {
+        if (disabled)
+          return
+        onClick()
+      }}>
+      {icon}
+      <div className='ml-1 text-[13px] text-gray-900'>{title}</div>
     </div>
   )
 })
diff --git a/web/app/components/base/prompt-editor/plugins/component-picker-block/variable-menu.tsx b/web/app/components/base/prompt-editor/plugins/component-picker-block/variable-menu.tsx
deleted file mode 100644
index fefd93cb0f..0000000000
--- a/web/app/components/base/prompt-editor/plugins/component-picker-block/variable-menu.tsx
+++ /dev/null
@@ -1,40 +0,0 @@
-import { memo } from 'react'
-import { VariableMenuItem } from './variable-option'
-
-type VariableMenuProps = {
-  startIndex: number
-  selectedIndex: number | null
-  options: any[]
-  onClick: (index: number, option: any) => void
-  onMouseEnter: (index: number, option: any) => void
-  queryString: string | null
-}
-const VariableMenu = ({
-  startIndex,
-  selectedIndex,
-  options,
-  onClick,
-  onMouseEnter,
-  queryString,
-}: VariableMenuProps) => {
-  return (
-    <div className='p-1'>
-      {
-        options.map((option, index: number) => (
-          <VariableMenuItem
-            startIndex={startIndex}
-            index={index}
-            isSelected={selectedIndex === index + startIndex}
-            onClick={onClick}
-            onMouseEnter={onMouseEnter}
-            key={option.key}
-            option={option}
-            queryString={queryString}
-          />
-        ))
-      }
-    </div>
-  )
-}
-
-export default memo(VariableMenu)
diff --git a/web/app/components/base/prompt-editor/plugins/component-picker-block/variable-option.tsx b/web/app/components/base/prompt-editor/plugins/component-picker-block/variable-option.tsx
index 76f76c8491..27a88ab665 100644
--- a/web/app/components/base/prompt-editor/plugins/component-picker-block/variable-option.tsx
+++ b/web/app/components/base/prompt-editor/plugins/component-picker-block/variable-option.tsx
@@ -1,60 +1,32 @@
 import { memo } from 'react'
-import { MenuOption } from '@lexical/react/LexicalTypeaheadMenuPlugin'
 
-export class VariableOption extends MenuOption {
+type VariableMenuItemProps = {
   title: string
   icon?: JSX.Element
   extraElement?: JSX.Element
-  keywords: Array<string>
-  keyboardShortcut?: string
-  onSelect: (queryString: string) => void
-
-  constructor(
-    title: string,
-    options: {
-      icon?: JSX.Element
-      extraElement?: JSX.Element
-      keywords?: Array<string>
-      keyboardShortcut?: string
-      onSelect: (queryString: string) => void
-    },
-  ) {
-    super(title)
-    this.title = title
-    this.keywords = options.keywords || []
-    this.icon = options.icon
-    this.extraElement = options.extraElement
-    this.keyboardShortcut = options.keyboardShortcut
-    this.onSelect = options.onSelect.bind(this)
-  }
-}
-
-type VariableMenuItemProps = {
-  startIndex: number
-  index: number
   isSelected: boolean
-  onClick: (index: number, option: VariableOption) => void
-  onMouseEnter: (index: number, option: VariableOption) => void
-  option: VariableOption
   queryString: string | null
+  onClick: () => void
+  onMouseEnter: () => void
+  setRefElement?: (element: HTMLDivElement) => void
 }
 export const VariableMenuItem = memo(({
-  startIndex,
-  index,
+  title,
+  icon,
+  extraElement,
   isSelected,
+  queryString,
   onClick,
   onMouseEnter,
-  option,
-  queryString,
+  setRefElement,
 }: VariableMenuItemProps) => {
-  const title = option.title
   let before = title
   let middle = ''
   let after = ''
 
   if (queryString) {
     const regex = new RegExp(queryString, 'i')
-    const match = regex.exec(option.title)
+    const match = regex.exec(title)
 
     if (match) {
       before = title.substring(0, match.index)
@@ -65,24 +37,23 @@ export const VariableMenuItem = memo(({
 
   return (
     <div
-      key={option.key}
       className={`
         flex items-center px-3 h-6 rounded-md hover:bg-primary-50 cursor-pointer
         ${isSelected && 'bg-primary-50'}
       `}
       tabIndex={-1}
-      ref={option.setRefElement}
-      onMouseEnter={() => onMouseEnter(index + startIndex, option)}
-      onClick={() => onClick(index + startIndex, option)}>
+      ref={setRefElement}
+      onMouseEnter={onMouseEnter}
+      onClick={onClick}>
       <div className='mr-2'>
-        {option.icon}
+        {icon}
       </div>
-      <div className='grow text-[13px] text-gray-900 truncate' title={option.title}>
+      <div className='grow text-[13px] text-gray-900 truncate' title={title}>
         {before}
         <span className='text-[#2970FF]'>{middle}</span>
         {after}
       </div>
-      {option.extraElement}
+      {extraElement}
     </div>
   )
 })