🧪 feat: Experimental: Enable Switching Endpoints Mid-Conversation (danny-avila#1483)

danny-avila · web-flow · commit 9e34c1cc11aa · 2024-01-03T19:17:42.000-05:00
* fix: load all existing conversation settings on refresh

* refactor(buildDefaultConvo): use `lastConversationSetup.endpointType` before `conversation.endpointType`

* refactor(TMessage/messageSchema): add `endpoint` field to messages to differentiate generation origin

* feat(useNewConvo): `keepLatestMessage` param to prevent reseting the `latestMessage` mid-conversation

* style(Settings): adjust height styling to allow more space in dialog for additional settings

* feat: Modular Chat: experimental setting to Enable switching Endpoints mid-conversation

* fix(ChatRoute): fix potential parsing issue with tPresetSchema
diff --git a/api/app/clients/BaseClient.js b/api/app/clients/BaseClient.js
@@ -516,7 +516,7 @@ class BaseClient {
   }
 
   async saveMessageToDatabase(message, endpointOptions, user = null) {
-    await saveMessage({ ...message, user, unfinished: false });
+    await saveMessage({ ...message, endpoint: this.options.endpoint, user, unfinished: false });
     await saveConvo(user, {
       conversationId: message.conversationId,
       endpoint: this.options.endpoint,
diff --git a/api/models/Message.js b/api/models/Message.js
@@ -9,6 +9,7 @@ module.exports = {
 
   async saveMessage({
     user,
+    endpoint,
     messageId,
     newMessageId,
     conversationId,
@@ -34,6 +35,7 @@ module.exports = {
 
       const update = {
         user,
+        endpoint,
         messageId: newMessageId || messageId,
         conversationId,
         parentMessageId,
diff --git a/api/models/schema/messageSchema.js b/api/models/schema/messageSchema.js
@@ -23,9 +23,11 @@ const messageSchema = mongoose.Schema(
       type: String,
       default: null,
     },
+    endpoint: {
+      type: String,
+    },
     conversationSignature: {
       type: String,
-      // required: true
     },
     clientId: {
       type: String,
@@ -35,7 +37,6 @@ const messageSchema = mongoose.Schema(
     },
     parentMessageId: {
       type: String,
-      // required: true
     },
     tokenCount: {
       type: Number,
diff --git a/api/server/controllers/AskController.js b/api/server/controllers/AskController.js
@@ -118,6 +118,8 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
       response = { ...response, ...metadata };
     }
 
+    response.endpoint = endpointOption.endpoint;
+
     if (client.options.attachments) {
       userMessage.files = client.options.attachments;
       delete userMessage.image_urls;
diff --git a/client/src/components/Chat/Menus/Endpoints/MenuItem.tsx b/client/src/components/Chat/Menus/Endpoints/MenuItem.tsx
@@ -1,11 +1,14 @@
 import { useState } from 'react';
 import { Settings } from 'lucide-react';
+import { useRecoilValue } from 'recoil';
 import { EModelEndpoint } from 'librechat-data-provider';
 import { useGetEndpointsQuery } from 'librechat-data-provider/react-query';
 import type { FC } from 'react';
-import { useLocalize, useUserKey } from '~/hooks';
+import type { TPreset } from 'librechat-data-provider';
+import { useLocalize, useUserKey, useDefaultConvo } from '~/hooks';
 import { SetKeyDialog } from '~/components/Input/SetKeyDialog';
 import { useChatContext } from '~/Providers';
+import store from '~/store';
 import { icons } from './Icons';
 import { cn } from '~/utils';
 
@@ -27,10 +30,12 @@ const MenuItem: FC<MenuItemProps> = ({
   userProvidesKey,
   ...rest
 }) => {
+  const modularChat = useRecoilValue(store.modularChat);
+  const [isDialogOpen, setDialogOpen] = useState(false);
   const { data: endpointsConfig } = useGetEndpointsQuery();
+  const { conversation, newConversation } = useChatContext();
+  const getDefaultConversation = useDefaultConvo();
 
-  const [isDialogOpen, setDialogOpen] = useState(false);
-  const { newConversation } = useChatContext();
   const { getExpiry } = useUserKey(endpoint);
   const localize = useLocalize();
   const expiryTime = getExpiry();
@@ -42,7 +47,22 @@ const MenuItem: FC<MenuItemProps> = ({
       if (!expiryTime) {
         setDialogOpen(true);
       }
-      newConversation({ template: { endpoint: newEndpoint, conversationId: 'new' } });
+      const template: Partial<TPreset> = { endpoint: newEndpoint, conversationId: 'new' };
+      const { conversationId } = conversation ?? {};
+      if (modularChat && conversationId && conversationId !== 'new') {
+        template.endpointType = endpointsConfig?.[newEndpoint]?.type;
+
+        const currentConvo = getDefaultConversation({
+          /* target endpointType is necessary to avoid endpoint mixing */
+          conversation: { ...(conversation ?? {}), endpointType: template.endpointType },
+          preset: template,
+        });
+
+        /* We don't reset the latest message, only when changing settings mid-converstion */
+        newConversation({ template: currentConvo, keepLatestMessage: true });
+        return;
+      }
+      newConversation({ template });
     }
   };
 
diff --git a/client/src/components/Nav/Settings.tsx b/client/src/components/Nav/Settings.tsx
@@ -1,9 +1,9 @@
 import * as Tabs from '@radix-ui/react-tabs';
+import type { TDialogProps } from '~/common';
 import { Dialog, DialogContent, DialogHeader, DialogTitle } from '~/components/ui';
 import { GearIcon, DataIcon, UserIcon } from '~/components/svg';
-import { useMediaQuery, useLocalize } from '~/hooks';
-import type { TDialogProps } from '~/common';
 import { General, Data, Account } from './SettingsTabs';
+import { useMediaQuery, useLocalize } from '~/hooks';
 import { cn } from '~/utils';
 
 export default function Settings({ open, onOpenChange }: TDialogProps) {
@@ -13,7 +13,7 @@ export default function Settings({ open, onOpenChange }: TDialogProps) {
   return (
     <Dialog open={open} onOpenChange={onOpenChange}>
       <DialogContent
-        className={cn('shadow-2xl dark:bg-gray-900 dark:text-white md:h-[373px] md:w-[680px]')}
+        className={cn('shadow-2xl dark:bg-gray-900 dark:text-white md:min-h-[373px] md:w-[680px]')}
         style={{ borderRadius: '12px' }}
       >
         <DialogHeader>
diff --git a/client/src/components/Nav/SettingsTabs/General/General.tsx b/client/src/components/Nav/SettingsTabs/General/General.tsx
@@ -12,9 +12,10 @@ import {
 } from '~/hooks';
 import type { TDangerButtonProps } from '~/common';
 import AutoScrollSwitch from './AutoScrollSwitch';
+import { Dropdown } from '~/components/ui';
 import DangerButton from '../DangerButton';
+import ModularChat from './ModularChat';
 import store from '~/store';
-import { Dropdown } from '~/components/ui';
 
 export const ThemeSelector = ({
   theme,
@@ -188,6 +189,9 @@ function General() {
         <div className="border-b pb-3 last-of-type:border-b-0 dark:border-gray-700">
           <AutoScrollSwitch />
         </div>
+        <div className="border-b pb-3 last-of-type:border-b-0 dark:border-gray-700">
+          <ModularChat />
+        </div>
       </div>
     </Tabs.Content>
   );
diff --git a/client/src/components/Nav/SettingsTabs/General/ModularChat.tsx b/client/src/components/Nav/SettingsTabs/General/ModularChat.tsx
@@ -0,0 +1,35 @@
+import { useRecoilState } from 'recoil';
+import { Switch } from '~/components/ui';
+import { useLocalize } from '~/hooks';
+import store from '~/store';
+
+export default function ModularChatSwitch({
+  onCheckedChange,
+}: {
+  onCheckedChange?: (value: boolean) => void;
+}) {
+  const [modularChat, setModularChat] = useRecoilState<boolean>(store.modularChat);
+  const localize = useLocalize();
+
+  const handleCheckedChange = (value: boolean) => {
+    setModularChat(value);
+    if (onCheckedChange) {
+      onCheckedChange(value);
+    }
+  };
+
+  return (
+    <div className="flex items-center justify-between">
+      <div>
+        {`[${localize('com_ui_experimental')}]`} {localize('com_nav_modular_chat')}{' '}
+      </div>
+      <Switch
+        id="modularChat"
+        checked={modularChat}
+        onCheckedChange={handleCheckedChange}
+        className="ml-4 mt-2"
+        data-testid="modularChat"
+      />
+    </div>
+  );
+}
diff --git a/client/src/hooks/Conversations/usePresets.ts b/client/src/hooks/Conversations/usePresets.ts
@@ -1,11 +1,11 @@
-import { QueryKeys, modularEndpoints } from 'librechat-data-provider';
-import { useCreatePresetMutation } from 'librechat-data-provider/react-query';
 import filenamify from 'filenamify';
-import { useCallback, useEffect, useRef } from 'react';
-import { useRecoilState, useSetRecoilState } from 'recoil';
 import exportFromJSON from 'export-from-json';
+import { useCallback, useEffect, useRef } from 'react';
 import { useQueryClient } from '@tanstack/react-query';
-import type { TPreset } from 'librechat-data-provider';
+import { QueryKeys, modularEndpoints } from 'librechat-data-provider';
+import { useRecoilState, useSetRecoilState, useRecoilValue } from 'recoil';
+import { useCreatePresetMutation } from 'librechat-data-provider/react-query';
+import type { TPreset, TEndpointsConfig } from 'librechat-data-provider';
 import {
   useUpdatePresetMutation,
   useDeletePresetMutation,
@@ -27,6 +27,7 @@ export default function usePresets() {
   const { showToast } = useToastContext();
   const { user, isAuthenticated } = useAuthContext();
 
+  const modularChat = useRecoilValue(store.modularChat);
   const [_defaultPreset, setDefaultPreset] = useRecoilState(store.defaultPreset);
   const setPresetModalVisible = useSetRecoilState(store.presetModalVisible);
   const { preset, conversation, newConversation, setPreset } = useChatContext();
@@ -159,14 +160,20 @@ export default function usePresets() {
       duration: 750,
     });
 
+    const endpointsConfig = queryClient.getQueryData<TEndpointsConfig>([QueryKeys.endpoints]);
+
+    const currentEndpointType = endpointsConfig?.[endpoint ?? '']?.type ?? '';
+    const endpointType = endpointsConfig?.[newPreset?.endpoint ?? '']?.type;
+
     if (
-      modularEndpoints.has(endpoint ?? '') &&
-      modularEndpoints.has(newPreset?.endpoint ?? '') &&
-      endpoint === newPreset?.endpoint
+      (modularEndpoints.has(endpoint ?? '') || modularEndpoints.has(currentEndpointType)) &&
+      (modularEndpoints.has(newPreset?.endpoint ?? '') || modularEndpoints.has(endpointType)) &&
+      (endpoint === newPreset?.endpoint || modularChat)
     ) {
       const currentConvo = getDefaultConversation({
-        conversation: conversation ?? {},
-        preset: newPreset,
+        /* target endpointType is necessary to avoid endpoint mixing */
+        conversation: { ...(conversation ?? {}), endpointType },
+        preset: { ...newPreset, endpointType },
       });
 
       /* We don't reset the latest message, only when changing settings mid-converstion */
diff --git a/client/src/hooks/useChatHelpers.ts b/client/src/hooks/useChatHelpers.ts
@@ -225,6 +225,7 @@ export default function useChatHelpers(index = 0, paramId: string | undefined) {
     const initialResponse: TMessage = {
       sender: responseSender,
       text: responseText,
+      endpoint: endpoint ?? '',
       parentMessageId: isRegenerate ? messageId : fakeMessageId,
       messageId: responseMessageId ?? `${isRegenerate ? messageId : fakeMessageId}_`,
       conversationId,
diff --git a/client/src/hooks/useNewConvo.ts b/client/src/hooks/useNewConvo.ts
@@ -46,6 +46,7 @@ const useNewConvo = (index = 0) => {
         preset: TPreset | null = null,
         modelsData?: TModelsConfig,
         buildDefault?: boolean,
+        keepLatestMessage?: boolean,
       ) => {
         const modelsConfig = modelsData ?? snapshot.getLoadable(store.modelsConfig).contents;
         const { endpoint = null } = conversation;
@@ -84,7 +85,9 @@ const useNewConvo = (index = 0) => {
         setStorage(conversation);
         setConversation(conversation);
         setSubmission({} as TSubmission);
-        resetLatestMessage();
+        if (!keepLatestMessage) {
+          resetLatestMessage();
+        }
 
         if (conversation.conversationId === 'new' && !modelsData) {
           navigate('new');
@@ -99,11 +102,13 @@ const useNewConvo = (index = 0) => {
       preset,
       modelsData,
       buildDefault = true,
+      keepLatestMessage = false,
     }: {
       template?: Partial<TConversation>;
       preset?: TPreset;
       modelsData?: TModelsConfig;
       buildDefault?: boolean;
+      keepLatestMessage?: boolean;
     } = {}) => {
       const conversation = {
         conversationId: 'new',
@@ -130,7 +135,7 @@ const useNewConvo = (index = 0) => {
         }
       }
 
-      switchToConversation(conversation, preset, modelsData, buildDefault);
+      switchToConversation(conversation, preset, modelsData, buildDefault, keepLatestMessage);
     },
     [switchToConversation, files, mutateAsync, setFiles],
   );
diff --git a/client/src/localization/languages/Eng.tsx b/client/src/localization/languages/Eng.tsx
@@ -15,6 +15,7 @@ export default {
   com_ui_limitation_harmful_biased:
     'May occasionally produce harmful instructions or biased content',
   com_ui_limitation_limited_2021: 'Limited knowledge of world and events after 2021',
+  com_ui_experimental: 'Experimental',
   com_ui_input: 'Input',
   com_ui_close: 'Close',
   com_ui_model: 'Model',
@@ -257,6 +258,7 @@ export default {
     'Make sure to click \'Create and Continue\' to give at least the \'Vertex AI User\' role. Lastly, create a JSON key to import here.',
   com_nav_welcome_message: 'How can I help you today?',
   com_nav_auto_scroll: 'Auto-scroll to Newest on Open',
+  com_nav_modular_chat: 'Enable switching Endpoints mid-conversation',
   com_nav_profile_picture: 'Profile Picture',
   com_nav_change_picture: 'Change picture',
   com_nav_plugin_store: 'Plugin store',
diff --git a/client/src/routes/ChatRoute.tsx b/client/src/routes/ChatRoute.tsx
@@ -6,6 +6,7 @@ import {
   useGetModelsQuery,
   useGetEndpointsQuery,
 } from 'librechat-data-provider/react-query';
+import { TPreset } from 'librechat-data-provider';
 import { useNewConvo, useConfigOverride } from '~/hooks';
 import ChatView from '~/components/Chat/ChatView';
 import useAuthRedirect from './useAuthRedirect';
@@ -45,6 +46,8 @@ export default function ChatRoute() {
     ) {
       newConversation({
         template: initialConvoQuery.data,
+        /* this is necessary to load all existing settings */
+        preset: initialConvoQuery.data as TPreset,
         modelsData: modelsQuery.data,
       });
       hasSetConversation.current = true;
diff --git a/client/src/store/settings.ts b/client/src/store/settings.ts
@@ -50,6 +50,25 @@ const autoScroll = atom<boolean>({
   ] as const,
 });
 
+const modularChat = atom<boolean>({
+  key: 'modularChat',
+  default: localStorage.getItem('modularChat') === 'true',
+  effects: [
+    ({ setSelf, onSet }) => {
+      const savedValue = localStorage.getItem('modularChat');
+      if (savedValue != null) {
+        setSelf(savedValue === 'true');
+      }
+
+      onSet((newValue: unknown) => {
+        if (typeof newValue === 'boolean') {
+          localStorage.setItem('modularChat', newValue.toString());
+        }
+      });
+    },
+  ] as const,
+});
+
 export default {
   abortScroll,
   optionSettings,
@@ -58,4 +77,5 @@ export default {
   showBingToneSetting,
   showPopover,
   autoScroll,
+  modularChat,
 };
diff --git a/client/src/utils/buildDefaultConvo.ts b/client/src/utils/buildDefaultConvo.ts
@@ -15,7 +15,7 @@ const buildDefaultConvo = ({
 }) => {
   const { lastSelectedModel, lastSelectedTools, lastBingSettings } = getLocalStorageItems();
   const { jailbreak, toneStyle } = lastBingSettings;
-  const { endpointType } = conversation;
+  const endpointType = lastConversationSetup?.endpointType ?? conversation?.endpointType;
 
   if (!endpoint) {
     return {
diff --git a/packages/data-provider/src/schemas.ts b/packages/data-provider/src/schemas.ts
@@ -106,6 +106,7 @@ export const tAgentOptionsSchema = z.object({
 
 export const tMessageSchema = z.object({
   messageId: z.string(),
+  endpoint: z.string().optional(),
   clientId: z.string().nullable().optional(),
   conversationId: z.string().nullable(),
   parentMessageId: z.string().nullable(),

Original file line number	Diff line number	Diff line change
`@@ -516,7 +516,7 @@ class BaseClient {`
`516`	`516`	`}`
`517`	`517`
`518`	`518`	`async saveMessageToDatabase(message, endpointOptions, user = null) {`
`519`		`- await saveMessage({ ...message, user, unfinished: false });`
	`519`	`+ await saveMessage({ ...message, endpoint: this.options.endpoint, user, unfinished: false });`
`520`	`520`	`await saveConvo(user, {`
`521`	`521`	`conversationId: message.conversationId,`
`522`	`522`	`endpoint: this.options.endpoint,`
Original file line number	Diff line number	Diff line change
`@@ -118,6 +118,8 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {`
`118`	`118`	`response = { ...response, ...metadata };`
`119`	`119`	`}`
`120`	`120`
	`121`	`+ response.endpoint = endpointOption.endpoint;`
	`122`	`+`
`121`	`123`	`if (client.options.attachments) {`
`122`	`124`	`userMessage.files = client.options.attachments;`
`123`	`125`	`delete userMessage.image_urls;`