From 4b92e84d122dcbbb68b969b75ea3472730ecbe44 Mon Sep 17 00:00:00 2001
From: Tung Pham <phsontung@gmail.com>
Date: Fri, 5 Apr 2024 08:20:34 +0700
Subject: [PATCH] Add support for change default translation agent

---
 locales/en-US/setting.json                    |  8 ++
 .../settings/common/features/SystemAgent.tsx  | 98 +++++++++++++++++++
 src/const/settings/index.ts                   | 10 ++
 src/locales/default/setting.ts                |  8 ++
 src/store/chat/slices/enchance/action.ts      | 16 ++-
 .../slices/settings/actions/general.test.ts   | 22 +++++
 .../user/slices/settings/actions/general.ts   | 11 +++
 .../slices/settings/selectors/settings.ts     | 11 ++-
 src/types/settings/index.ts                   |  3 +
 src/types/settings/systemAgent.ts             |  8 ++
 10 files changed, 192 insertions(+), 3 deletions(-)
 create mode 100644 src/app/(main)/settings/common/features/SystemAgent.tsx
 create mode 100644 src/types/settings/systemAgent.ts

diff --git a/locales/en-US/setting.json b/locales/en-US/setting.json
index 5b6bd98f7a9e..e302932757d8 100644
--- a/locales/en-US/setting.json
+++ b/locales/en-US/setting.json
@@ -355,6 +355,14 @@
       "title": "WebRTC Sync"
     }
   },
+  "systemAgent": {
+    "title": "System Agents",
+    "translation": {
+      "title": "Translation Settings",
+      "label": "Translation Agent",
+      "modelDesc": "Specific model for translate message"
+    }
+  },
   "tab": {
     "about": "About",
     "agent": "Default Agent",
diff --git a/src/app/(main)/settings/common/features/SystemAgent.tsx b/src/app/(main)/settings/common/features/SystemAgent.tsx
new file mode 100644
index 000000000000..13fad59b0673
--- /dev/null
+++ b/src/app/(main)/settings/common/features/SystemAgent.tsx
@@ -0,0 +1,98 @@
+import { Form, type ItemGroup } from '@lobehub/ui';
+import { Form as AntForm, Select, SelectProps } from 'antd';
+import { createStyles } from 'antd-style';
+import isEqual from 'fast-deep-equal';
+import { MessageSquareMore } from 'lucide-react';
+import { memo, useMemo } from 'react';
+import { useTranslation } from 'react-i18next';
+
+import { ModelItemRender, ProviderItemRender } from '@/components/ModelSelect';
+import { FORM_STYLE } from '@/const/layoutTokens';
+import { useUserStore } from '@/store/user';
+import { modelProviderSelectors, settingsSelectors } from '@/store/user/selectors';
+import { ModelProviderCard } from '@/types/llm';
+
+import { useSyncSettings } from '../../hooks/useSyncSettings';
+
+const SYSTEM_AGENT_SETTING_KEY = 'systemAgent';
+
+type SettingItemGroup = ItemGroup;
+
+const useStyles = createStyles(({ css, prefixCls }) => ({
+  select: css`
+    .${prefixCls}-select-dropdown .${prefixCls}-select-item-option-grouped {
+      padding-inline-start: 12px;
+    }
+  `,
+}));
+interface ModelOption {
+  label: any;
+  provider: string;
+  value: string;
+}
+
+const SystemAgent = memo(() => {
+  const { t } = useTranslation('setting');
+  const [form] = AntForm.useForm();
+
+  const settings = useUserStore(settingsSelectors.currentSettings, isEqual);
+  const [setTranslationSystemAgent] = useUserStore((s) => [s.setTranslationSystemAgent]);
+
+  const select = useUserStore(modelProviderSelectors.modelProviderListForModelSelect, isEqual);
+
+  const { styles } = useStyles();
+
+  const enabledList = select.filter((s) => s.enabled);
+
+  const translationOptions = useMemo<SelectProps['options']>(() => {
+    const getChatModels = (provider: ModelProviderCard) =>
+      provider.chatModels
+        .filter((c) => c.enabled)
+        .map((model) => ({
+          label: <ModelItemRender {...model} />,
+          provider: provider.id,
+          value: model.id,
+        }));
+
+    if (enabledList.length === 1) {
+      const provider = enabledList[0];
+
+      return getChatModels(provider);
+    }
+
+    return enabledList.map((provider) => ({
+      label: <ProviderItemRender provider={provider.id} />,
+      options: getChatModels(provider),
+    }));
+  }, [enabledList]);
+
+  const systemAgentSettings: SettingItemGroup = {
+    children: [
+      {
+        children: (
+          <Select
+            className={styles.select}
+            onChange={(model, option) =>
+              setTranslationSystemAgent((option as unknown as ModelOption).provider, model)
+            }
+            options={translationOptions}
+            popupMatchSelectWidth={false}
+          />
+        ),
+        desc: t('systemAgent.translation.modelDesc'),
+        label: t('systemAgent.translation.label'),
+        name: [SYSTEM_AGENT_SETTING_KEY, 'translation', 'model'],
+      },
+    ],
+    icon: MessageSquareMore,
+    title: t('systemAgent.title'),
+  };
+
+  useSyncSettings(form);
+
+  return (
+    <Form form={form} initialValues={settings} items={[systemAgentSettings]} {...FORM_STYLE} />
+  );
+});
+
+export default SystemAgent;
diff --git a/src/const/settings/index.ts b/src/const/settings/index.ts
index ee2330de4050..3ee13d15c4a0 100644
--- a/src/const/settings/index.ts
+++ b/src/const/settings/index.ts
@@ -24,7 +24,9 @@ import {
   GlobalLLMConfig,
   GlobalSettings,
   GlobalSyncSettings,
+  GlobalSystemAgentConfig,
   GlobalTTSConfig,
+  GlobalTranslationConfig,
 } from '@/types/settings';
 
 export const DEFAULT_BASE_SETTINGS: GlobalBaseSettings = {
@@ -158,6 +160,13 @@ export const DEFAULT_TTS_CONFIG: GlobalTTSConfig = {
   sttServer: 'openai',
 };
 
+export const DEFAULT_TRANSLATION_CONFIG: GlobalTranslationConfig = {
+  model: DEFAULT_AGENT_CONFIG.model,
+  provider: DEFAULT_AGENT_CONFIG.provider || ModelProvider.OpenAI,
+};
+export const DEFAULT_SYSTEM_AGENT_CONFIG: GlobalSystemAgentConfig = {
+  translation: DEFAULT_TRANSLATION_CONFIG,
+};
 export const DEFAULT_TOOL_CONFIG = {
   dalle: {
     autoGenerate: false,
@@ -172,6 +181,7 @@ export const DEFAULT_SETTINGS: GlobalSettings = {
   defaultAgent: DEFAULT_AGENT,
   languageModel: DEFAULT_LLM_CONFIG,
   sync: DEFAULT_SYNC_CONFIG,
+  systemAgent: DEFAULT_SYSTEM_AGENT_CONFIG,
   tool: DEFAULT_TOOL_CONFIG,
   tts: DEFAULT_TTS_CONFIG,
   ...DEFAULT_BASE_SETTINGS,
diff --git a/src/locales/default/setting.ts b/src/locales/default/setting.ts
index 3554db126135..3526716540ee 100644
--- a/src/locales/default/setting.ts
+++ b/src/locales/default/setting.ts
@@ -357,6 +357,14 @@ export default {
       title: 'WebRTC 同步',
     },
   },
+  systemAgent: {
+    title: 'System Agents',
+    translation: {
+      label: 'Translation model',
+      modelDesc: 'Specific model for translate message',
+      title: 'Translation Settings',
+    },
+  },
   tab: {
     about: '关于',
     agent: '默认助手',
diff --git a/src/store/chat/slices/enchance/action.ts b/src/store/chat/slices/enchance/action.ts
index a93d91e0de44..65659b84ae68 100644
--- a/src/store/chat/slices/enchance/action.ts
+++ b/src/store/chat/slices/enchance/action.ts
@@ -9,7 +9,11 @@ import { chatService } from '@/services/chat';
 import { messageService } from '@/services/message';
 import { chatSelectors } from '@/store/chat/selectors';
 import { ChatStore } from '@/store/chat/store';
+import { useUserStore } from '@/store/user';
+import { settingsSelectors } from '@/store/user/selectors';
 import { ChatTTS, ChatTranslate } from '@/types/message';
+import { GlobalTranslationConfig } from '@/types/settings';
+import { merge } from '@/utils/merge';
 import { setNamespace } from '@/utils/storeDebug';
 
 const n = setNamespace('enhance');
@@ -21,6 +25,7 @@ export interface ChatEnhanceAction {
   clearTTS: (id: string) => Promise<void>;
   clearTranslate: (id: string) => Promise<void>;
   getCurrentTracePayload: (data: Partial<TracePayload>) => TracePayload;
+  getCurrentTranslationSetting: () => GlobalTranslationConfig;
   translateMessage: (id: string, targetLang: string) => Promise<void>;
   ttsMessage: (
     id: string,
@@ -48,12 +53,19 @@ export const chatEnhance: StateCreator<
     topicId: get().activeTopicId,
     ...data,
   }),
+  getCurrentTranslationSetting: () => {
+    return settingsSelectors.currentSystemAgent(useUserStore.getState()).translation;
+  },
+
   translateMessage: async (id, targetLang) => {
     const { toggleChatLoading, updateMessageTranslate, dispatchMessage } = get();
 
     const message = chatSelectors.getMessageById(id)(get());
     if (!message) return;
 
+    // Get current agent for translation
+    const translationSetting = get().getCurrentTranslationSetting();
+
     // create translate extra
     await updateMessageTranslate(id, { content: '', from: '', to: targetLang });
 
@@ -65,7 +77,7 @@ export const chatEnhance: StateCreator<
     // detect from language
     chatService
       .fetchPresetTaskResult({
-        params: chainLangDetect(message.content),
+        params: merge(translationSetting, chainLangDetect(message.content)),
         trace: get().getCurrentTracePayload({ traceName: TraceNameMap.LanguageDetect }),
       })
       .then(async (data) => {
@@ -87,7 +99,7 @@ export const chatEnhance: StateCreator<
           }),
         });
       },
-      params: chainTranslate(message.content, targetLang),
+      params: merge(translationSetting, chainTranslate(message.content, targetLang)),
       trace: get().getCurrentTracePayload({ traceName: TraceNameMap.Translator }),
     });
 
diff --git a/src/store/user/slices/settings/actions/general.test.ts b/src/store/user/slices/settings/actions/general.test.ts
index fbe739410a7d..0565abbe1889 100644
--- a/src/store/user/slices/settings/actions/general.test.ts
+++ b/src/store/user/slices/settings/actions/general.test.ts
@@ -112,4 +112,26 @@ describe('SettingsAction', () => {
       expect(userService.updateUserSettings).toHaveBeenCalledWith({ defaultAgent: updatedAgent });
     });
   });
+
+  describe('setTranslationSystemAgent', () => {
+    it('should set partial settings', async () => {
+      const { result } = renderHook(() => useUserStore());
+      const systemAgentSettings: Partial<GlobalSettings> = {
+        systemAgent: {
+          translation: {
+            model: 'testmodel',
+            provider: 'provider',
+          },
+        },
+      };
+
+      // Perform the action
+      await act(async () => {
+        await result.current.setTranslationSystemAgent('provider', 'testmodel');
+      });
+
+      // Assert that updateUserSettings was called with the correct settings
+      expect(userService.updateUserSettings).toHaveBeenCalledWith(systemAgentSettings);
+    });
+  });
 });
diff --git a/src/store/user/slices/settings/actions/general.ts b/src/store/user/slices/settings/actions/general.ts
index 7a53916abfdb..a2a657a4d8a3 100644
--- a/src/store/user/slices/settings/actions/general.ts
+++ b/src/store/user/slices/settings/actions/general.ts
@@ -17,6 +17,7 @@ export interface GeneralSettingsAction {
   resetSettings: () => Promise<void>;
   setSettings: (settings: DeepPartial<GlobalSettings>) => Promise<void>;
   switchLocale: (locale: LocaleMode) => Promise<void>;
+  setTranslationSystemAgent: (provider: string, model: string) => Promise<void>;
   switchThemeMode: (themeMode: ThemeMode) => Promise<void>;
   updateDefaultAgent: (agent: DeepPartial<LobeAgentSettings>) => Promise<void>;
 }
@@ -55,6 +56,16 @@ export const generalSettingsSlice: StateCreator<
 
     switchLang(locale);
   },
+  setTranslationSystemAgent: async (provider, model) => {
+    await get().setSettings({
+      systemAgent: {
+        translation: {
+          model: model,
+          provider: provider,
+        },
+      },
+    });
+  },
   switchThemeMode: async (themeMode) => {
     await get().setSettings({ themeMode });
   },
diff --git a/src/store/user/slices/settings/selectors/settings.ts b/src/store/user/slices/settings/selectors/settings.ts
index cbb6f310f5b3..bdcb832665c2 100644
--- a/src/store/user/slices/settings/selectors/settings.ts
+++ b/src/store/user/slices/settings/selectors/settings.ts
@@ -1,6 +1,11 @@
 import { DEFAULT_LANG } from '@/const/locale';
 import { DEFAULT_AGENT_META } from '@/const/meta';
-import { DEFAULT_AGENT, DEFAULT_AGENT_CONFIG, DEFAULT_TTS_CONFIG } from '@/const/settings';
+import {
+  DEFAULT_AGENT,
+  DEFAULT_AGENT_CONFIG,
+  DEFAULT_SYSTEM_AGENT_CONFIG,
+  DEFAULT_TTS_CONFIG,
+} from '@/const/settings';
 import { Locales } from '@/locales/resources';
 import { GeneralModelProviderConfig, GlobalLLMProviderKey, GlobalSettings } from '@/types/settings';
 import { isOnServerSide } from '@/utils/env';
@@ -48,9 +53,13 @@ const currentLanguage = (s: UserStore) => {
 const dalleConfig = (s: UserStore) => currentSettings(s).tool?.dalle || {};
 const isDalleAutoGenerating = (s: UserStore) => currentSettings(s).tool?.dalle?.autoGenerate;
 
+const currentSystemAgent = (s: UserStore) =>
+  merge(DEFAULT_SYSTEM_AGENT_CONFIG, currentSettings(s).systemAgent);
+
 export const settingsSelectors = {
   currentLanguage,
   currentSettings,
+  currentSystemAgent,
   currentTTS,
   dalleConfig,
   defaultAgent,
diff --git a/src/types/settings/index.ts b/src/types/settings/index.ts
index cff683aff51d..c7c03c73cafd 100644
--- a/src/types/settings/index.ts
+++ b/src/types/settings/index.ts
@@ -3,6 +3,7 @@ import type { LobeAgentSession } from '@/types/session';
 import { GlobalBaseSettings } from './base';
 import { GlobalLLMConfig } from './modelProvider';
 import { GlobalSyncSettings } from './sync';
+import { GlobalSystemAgentConfig } from './systemAgent';
 import { GlobalTTSConfig } from './tts';
 
 export type GlobalDefaultAgent = Pick<LobeAgentSession, 'config' | 'meta'>;
@@ -10,6 +11,7 @@ export type GlobalDefaultAgent = Pick<LobeAgentSession, 'config' | 'meta'>;
 export * from './base';
 export * from './modelProvider';
 export * from './sync';
+export * from './systemAgent';
 export * from './tts';
 
 export interface GlobalTool {
@@ -25,6 +27,7 @@ export interface GlobalSettings extends GlobalBaseSettings {
   defaultAgent: GlobalDefaultAgent;
   languageModel: GlobalLLMConfig;
   sync: GlobalSyncSettings;
+  systemAgent: GlobalSystemAgentConfig;
   tool: GlobalTool;
   tts: GlobalTTSConfig;
 }
diff --git a/src/types/settings/systemAgent.ts b/src/types/settings/systemAgent.ts
new file mode 100644
index 000000000000..90b86273b389
--- /dev/null
+++ b/src/types/settings/systemAgent.ts
@@ -0,0 +1,8 @@
+export interface GlobalTranslationConfig {
+  model: string;
+  provider: string;
+}
+
+export interface GlobalSystemAgentConfig {
+  translation: GlobalTranslationConfig;
+}