♻️ refactor: refactor agent store data (lobehub#2690)

* ♻️ refactor: refactor agent store data * 🎨 refactor: improve code
bentwnghk · May 28, 2024 · e201937 · e201937
1 parent d420de3
commit e201937
Show file tree

Hide file tree

Showing 14 changed files with 172 additions and 167 deletions.
diff --git a/src/app/(main)/chat/@session/features/SessionListContent/List/Item/index.tsx b/src/app/(main)/chat/@session/features/SessionListContent/List/Item/index.tsx
@@ -4,6 +4,7 @@ import { shallow } from 'zustand/shallow';
 
 import ModelTag from '@/components/ModelTag';
 import { useAgentStore } from '@/store/agent';
+import { agentSelectors } from '@/store/agent/selectors';
 import { useChatStore } from '@/store/chat';
 import { chatSelectors } from '@/store/chat/selectors';
 import { useSessionStore } from '@/store/session';
@@ -21,7 +22,7 @@ interface SessionItemProps {
 const SessionItem = memo<SessionItemProps>(({ id }) => {
   const [open, setOpen] = useState(false);
   const [createGroupModalOpen, setCreateGroupModalOpen] = useState(false);
-  const [defaultModel] = useAgentStore((s) => [s.defaultAgentConfig.model]);
+  const [defaultModel] = useAgentStore((s) => [agentSelectors.inboxAgentModel(s)]);
 
   const [active] = useSessionStore((s) => [s.activeId === id]);
   const [loading] = useChatStore((s) => [chatSelectors.isAIGenerating(s) && id === s.activeId]);

diff --git a/src/app/(main)/chat/settings/features/EditPage.tsx b/src/app/(main)/chat/settings/features/EditPage.tsx
@@ -16,10 +16,7 @@ const EditPage = memo(() => {
   const id = useSessionStore((s) => s.activeId);
   const config = useAgentStore(agentSelectors.currentAgentConfig, isEqual);
   const meta = useSessionStore(sessionMetaSelectors.currentAgentMeta, isEqual);
-  const [updateAgentConfig, updateAgentChatConfig] = useAgentStore((s) => [
-    s.updateAgentConfig,
-    s.updateAgentChatConfig,
-  ]);
+  const [updateAgentConfig] = useAgentStore((s) => [s.updateAgentConfig]);
 
   const [updateAgentMeta, title] = useSessionStore((s) => [
     s.updateSessionMeta,
@@ -33,7 +30,6 @@ const EditPage = memo(() => {
         config={config}
         id={id}
         meta={meta}
-        onChatConfigChange={updateAgentChatConfig}
         onConfigChange={updateAgentConfig}
         onMetaChange={updateAgentMeta}
       />

diff --git a/src/app/@modal/chat/(.)settings/modal/layout.tsx b/src/app/@modal/chat/(.)settings/modal/layout.tsx
@@ -28,10 +28,7 @@ const Layout = memo<PropsWithChildren>(({ children }) => {
   const id = useSessionStore((s) => s.activeId);
   const config = useAgentStore(agentSelectors.currentAgentConfig, isEqual);
   const meta = useSessionStore(sessionMetaSelectors.currentAgentMeta, isEqual);
-  const [updateAgentConfig, updateAgentChatConfig] = useAgentStore((s) => [
-    s.updateAgentConfig,
-    s.updateAgentChatConfig,
-  ]);
+  const [updateAgentConfig] = useAgentStore((s) => [s.updateAgentConfig]);
   const [updateAgentMeta] = useSessionStore((s) => [
     s.updateSessionMeta,
     sessionMetaSelectors.currentAgentTitle(s),
@@ -49,7 +46,6 @@ const Layout = memo<PropsWithChildren>(({ children }) => {
           config={config}
           id={id}
           meta={meta}
-          onChatConfigChange={updateAgentChatConfig}
           onConfigChange={updateAgentConfig}
           onMetaChange={updateAgentMeta}
         />

diff --git a/src/features/AgentSetting/StoreUpdater.tsx b/src/features/AgentSetting/StoreUpdater.tsx
@@ -6,19 +6,18 @@ import { createStoreUpdater } from 'zustand-utils';
 import { State, useStoreApi } from './store';
 
 export type StoreUpdaterProps = Partial<
-  Pick<State, 'onMetaChange' | 'onChatConfigChange' | 'onConfigChange' | 'meta' | 'config' | 'id'>
+  Pick<State, 'onMetaChange' | 'onConfigChange' | 'meta' | 'config' | 'id'>
 >;
 
 const StoreUpdater = memo<StoreUpdaterProps>(
-  ({ onConfigChange, onChatConfigChange, id, onMetaChange, meta, config }) => {
+  ({ onConfigChange, id, onMetaChange, meta, config }) => {
     const storeApi = useStoreApi();
     const useStoreUpdater = createStoreUpdater(storeApi);
 
     useStoreUpdater('meta', meta);
     useStoreUpdater('config', config);
     useStoreUpdater('onConfigChange', onConfigChange);
     useStoreUpdater('onMetaChange', onMetaChange);
-    useStoreUpdater('onChatConfigChange', onChatConfigChange);
     useStoreUpdater('id', id);
 
     return null;

diff --git a/src/features/AgentSetting/store/action.ts b/src/features/AgentSetting/store/action.ts
@@ -1,3 +1,4 @@
+import { DeepPartial } from 'utility-types';
 import { StateCreator } from 'zustand/vanilla';
 
 import { chainPickEmoji } from '@/chains/pickEmoji';
@@ -50,7 +51,7 @@ export interface Action {
   resetAgentConfig: () => void;
   resetAgentMeta: () => void;
 
-  setAgentConfig: (config: Partial<LobeAgentConfig>) => void;
+  setAgentConfig: (config: DeepPartial<LobeAgentConfig>) => void;
   setAgentMeta: (meta: Partial<MetaData>) => void;
   setChatConfig: (config: Partial<LobeAgentChatConfig>) => void;
 
@@ -245,11 +246,7 @@ export const store: StateCreator<Store, [['zustand/devtools', never]]> = (set, g
     get().dispatchMeta({ type: 'update', value: meta });
   },
   setChatConfig: (config) => {
-    const nextConfig = { ...get().config.chatConfig, ...config };
-
-    set({ config: { ...get().config, chatConfig: nextConfig } }, false, 'updateChatConfig');
-
-    get().onChatConfigChange?.(nextConfig);
+    get().setAgentConfig({ chatConfig: config });
   },
 
   streamUpdateMetaArray: (key: keyof MetaData) => {

diff --git a/src/features/AgentSetting/store/initialState.ts b/src/features/AgentSetting/store/initialState.ts
@@ -1,6 +1,6 @@
 import { DEFAULT_AGENT_META } from '@/const/meta';
 import { DEFAULT_AGENT_CONFIG } from '@/const/settings';
-import { LobeAgentChatConfig, LobeAgentConfig } from '@/types/agent';
+import { LobeAgentConfig } from '@/types/agent';
 import { MetaData } from '@/types/meta';
 
 export interface State {
@@ -9,7 +9,6 @@ export interface State {
   id?: string;
   meta: MetaData;
 
-  onChatConfigChange?: (config: LobeAgentChatConfig) => void;
   onConfigChange?: (config: LobeAgentConfig) => void;
   onMetaChange?: (meta: MetaData) => void;
 }

diff --git a/src/layout/GlobalProvider/StoreInitialization.tsx b/src/layout/GlobalProvider/StoreInitialization.tsx
@@ -26,10 +26,13 @@ const StoreInitialization = memo(() => {
 
   const useInitSystemStatus = useGlobalStore((s) => s.useInitSystemStatus);
 
-  const useFetchDefaultAgentConfig = useAgentStore((s) => s.useFetchDefaultAgentConfig);
+  const useInitAgentStore = useAgentStore((s) => s.useInitAgentStore);
+
   // init the system preference
   useInitSystemStatus();
-  useFetchDefaultAgentConfig();
+
+  // init inbox agent and default agent config
+  useInitAgentStore(serverConfig.defaultAgent?.config);
 
   useInitUserState(isLogin, serverConfig, {
     onSuccess: (state) => {
@@ -55,21 +58,21 @@ const StoreInitialization = memo(() => {
     importSettings(searchParam);
   }, [searchParam]);
 
-  useEffect(() => {
-    router.prefetch('/chat');
-    router.prefetch('/market');
-
-    if (mobile) {
-      router.prefetch('/me');
-      router.prefetch('/chat/settings');
-      router.prefetch('/settings/common');
-      router.prefetch('/settings/agent');
-      router.prefetch('/settings/sync');
-    } else {
-      router.prefetch('/chat/settings/modal');
-      router.prefetch('/settings/modal');
-    }
-  }, [router, mobile]);
+  // useEffect(() => {
+  //   router.prefetch('/chat');
+  //   router.prefetch('/market');
+  //
+  //   if (mobile) {
+  //     router.prefetch('/me');
+  //     router.prefetch('/chat/settings');
+  //     router.prefetch('/settings/common');
+  //     router.prefetch('/settings/agent');
+  //     router.prefetch('/settings/sync');
+  //   } else {
+  //     router.prefetch('/chat/settings/modal');
+  //     router.prefetch('/settings/modal');
+  //   }
+  // }, [router, mobile]);
 
   return null;
 });

diff --git a/src/store/agent/slices/chat/action.test.ts b/src/store/agent/slices/chat/action.test.ts
@@ -2,6 +2,7 @@ import { act, renderHook, waitFor } from '@testing-library/react';
 import { mutate } from 'swr';
 import { describe, expect, it, vi } from 'vitest';
 
+import { INBOX_SESSION_ID } from '@/const/session';
 import { DEFAULT_AGENT_CONFIG } from '@/const/settings';
 import { globalService } from '@/services/global';
 import { sessionService } from '@/services/session';
@@ -159,26 +160,29 @@ describe('AgentSlice', () => {
     it('should update agentConfig and isAgentConfigInit when data changes and isAgentConfigInit is false', async () => {
       const { result } = renderHook(() => useAgentStore());
 
-      act(() => {
-        result.current.isAgentConfigInit = false;
-        result.current.agentConfig = { model: 'gpt-3.5-turbo' };
-      });
+      // act(() => {
+      //   result.current.agentMap = {};
+      // });
 
       vi.spyOn(sessionService, 'getSessionConfig').mockResolvedValueOnce({ model: 'gpt-4' } as any);
 
       renderHook(() => result.current.useFetchAgentConfig('test-session-id'));
 
       await waitFor(() => {
-        expect(result.current.agentConfig).toEqual({ model: 'gpt-4' });
-        expect(result.current.isAgentConfigInit).toBe(true);
+        expect(result.current.agentMap['test-session-id']).toEqual({ model: 'gpt-4' });
+        // expect(result.current.isAgentConfigInit).toBe(true);
       });
     });
 
     it('should not update state when data is the same and isAgentConfigInit is true', async () => {
       const { result } = renderHook(() => useAgentStore());
 
       act(() => {
-        useAgentStore.setState({ isAgentConfigInit: true });
+        useAgentStore.setState({
+          agentMap: {
+            'test-session-id': { model: 'gpt-3.5-turbo' },
+          },
+        });
       });
 
       vi.spyOn(useSessionStore, 'setState');
@@ -189,27 +193,25 @@ describe('AgentSlice', () => {
       renderHook(() => result.current.useFetchAgentConfig('test-session-id'));
 
       await waitFor(() => {
-        expect(result.current.agentConfig).toEqual({ model: 'gpt-3.5-turbo' });
-        expect(result.current.isAgentConfigInit).toBe(true);
+        expect(result.current.agentMap['test-session-id']).toEqual({ model: 'gpt-3.5-turbo' });
 
         expect(useSessionStore.setState).not.toHaveBeenCalled();
       });
     });
   });
 
-  describe('useFetchDefaultAgentConfig', () => {
+  describe('useFetchInboxAgentConfig', () => {
     it('should merge DEFAULT_AGENT_CONFIG and update defaultAgentConfig and isDefaultAgentConfigInit on success', async () => {
       const { result } = renderHook(() => useAgentStore());
-      vi.spyOn(globalService, 'getDefaultAgentConfig').mockResolvedValue({ model: 'gemini-pro' });
+      vi.spyOn(sessionService, 'getSessionConfig').mockResolvedValue({
+        model: 'gemini-pro',
+      } as any);
 
-      renderHook(() => result.current.useFetchDefaultAgentConfig());
+      renderHook(() => result.current.useInitAgentStore());
 
       await waitFor(async () => {
-        expect(result.current.defaultAgentConfig).toEqual({
-          ...DEFAULT_AGENT_CONFIG,
-          model: 'gemini-pro',
-        });
-        expect(result.current.isDefaultAgentConfigInit).toBe(true);
+        expect(result.current.agentMap[INBOX_SESSION_ID]).toEqual({ model: 'gemini-pro' });
+        expect(result.current.isInboxAgentConfigInit).toBe(true);
       });
     });
 
@@ -218,11 +220,11 @@ describe('AgentSlice', () => {
 
       vi.spyOn(globalService, 'getDefaultAgentConfig').mockRejectedValueOnce(new Error());
 
-      renderHook(() => result.current.useFetchDefaultAgentConfig());
+      renderHook(() => result.current.useInitAgentStore());
 
       await waitFor(async () => {
-        expect(result.current.defaultAgentConfig).toEqual(DEFAULT_AGENT_CONFIG);
-        expect(result.current.isDefaultAgentConfigInit).toBe(false);
+        expect(result.current.agentMap[INBOX_SESSION_ID]).toBeUndefined();
+        expect(result.current.isInboxAgentConfigInit).toBe(false);
       });
     });
   });