Support config whether show mindmap

memfreeme · Nov 26, 2024 · 0fd79ea · 0fd79ea
1 parent 6d2d110
commit 0fd79ea
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 7 deletions.
diff --git a/frontend/components/search/search-bar.tsx b/frontend/components/search/search-bar.tsx
@@ -203,7 +203,7 @@ const SearchBar: React.FC<Props> = ({
         loading: () => <></>,
     });
 
-    const { isSearch, isShadcnUI, setIsSearch, setIsShadcnUI } = useUIStore();
+    const { isSearch, isShadcnUI, showMindMap, setIsSearch, setIsShadcnUI, setShowMindMap } = useUIStore();
 
     return (
         <div className="w-full text-center">
@@ -325,7 +325,7 @@ const SearchBar: React.FC<Props> = ({
                 </div>
             </div>
 
-            <div className="mt-4 grid grid-cols-1 md:grid-cols-3 gap-2">
+            <div className="mt-4 grid grid-cols-1 md:grid-cols-4 gap-2">
                 {showShadcnUI && (
                     <div className="flex items-center space-x-2 mb-1">
                         <Switch id="shadcn" checked={isShadcnUI} onCheckedChange={(checked) => setIsShadcnUI(checked)} />
@@ -334,6 +334,12 @@ const SearchBar: React.FC<Props> = ({
                 )}
                 {showModelSelection && <ModelSelection />}
                 {showSourceSelection && <SourceSelection />}
+                {showModelSelection && (
+                    <div className="flex items-center space-x-2 mb-1">
+                        <Switch id="mindmap" checked={showMindMap} onCheckedChange={(checked) => setShowMindMap(checked)} />
+                        <Label htmlFor="mindmap">Show MindMap</Label>
+                    </div>
+                )}
                 {showWebSearch && (
                     <div className="flex items-center space-x-2 mb-1">
                         <Switch id="search" checked={isSearch} onCheckedChange={(checked) => setIsSearch(checked)} />

diff --git a/frontend/components/search/search-message.tsx b/frontend/components/search/search-message.tsx
@@ -13,6 +13,7 @@ import ExpandableSection from '@/components/search/expandable-section';
 import MindMap from '@/components/search/mindmap';
 import { useTranslations } from 'next-intl';
 import UISection from '@/components/code/ui-section';
+import { useUIStore } from '@/lib/store';
 
 const SearchMessage = memo(
     (props: {
@@ -48,6 +49,7 @@ const SearchMessage = memo(
         }, [message.attachments, isUser, content]);
 
         const t = useTranslations('SearchMessage');
+        const { showMindMap } = useUIStore();
 
         return (
             <div className="flex flex-col w-full items-start space-y-6 pb-10">
@@ -62,11 +64,6 @@ const SearchMessage = memo(
                 {(images.length > 0 || !isLoading) && !isUser && !isReadOnly && (
                     <ActionButtons content={content} searchId={searchId} msgId={id} reload={reload} searchType={type} />
                 )}
-                {!isUser && content && !isLoading && type != SearchCategory.UI && (
-                    <ExpandableSection title={t('MindMap')} icon={Map} open={isReadOnly}>
-                        <MindMap value={content} />
-                    </ExpandableSection>
-                )}
                 {sources.length > 0 && (
                     <ExpandableSection title={t('Sources')} icon={TextSearchIcon}>
                         <div className="grid grid-cols-1 md:grid-cols-2 gap-4">
@@ -113,6 +110,12 @@ const SearchMessage = memo(
                     </div>
                 )}
 
+                {showMindMap && !isUser && content && !isLoading && type != SearchCategory.UI && (
+                    <ExpandableSection title={t('MindMap')} icon={Map}>
+                        <MindMap value={content} />
+                    </ExpandableSection>
+                )}
+
                 {isUser && <QuestionSection mesageId={id} content={content} isShared={isReadOnly} onContentChange={onSelect} reload={reload}></QuestionSection>}
                 {isUser && attachments && attachments.length > 0 && (
                     <div className="grid grid-cols-2 sm:grid-cols-3 md:grid-cols-4 gap-4 p-4">

diff --git a/frontend/lib/store.ts b/frontend/lib/store.ts
@@ -23,17 +23,21 @@ export const useProfileStore = create<ProfileState>()(
 type UIState = {
     isSearch: boolean;
     isShadcnUI: boolean;
+    showMindMap: boolean;
     setIsSearch: (isSearch: boolean) => void;
     setIsShadcnUI: (isShadcnUI: boolean) => void;
+    setShowMindMap: (showMindMap: boolean) => void;
 };
 
 export const useUIStore = create<UIState>()(
     persist(
         (set) => ({
             isSearch: true,
             isShadcnUI: true,
+            showMindMap: false,
             setIsSearch: (search: boolean) => set({ isSearch: search }),
             setIsShadcnUI: (shadcnUI: boolean) => set({ isShadcnUI: shadcnUI }),
+            setShowMindMap: (showMindMap: boolean) => set({ showMindMap }),
         }),
         {
             name: 'UI-config',