refactor: Encrypt & Expire User Provided Keys, feat: Rate Limiting (#874)

* docs: make_your_own.md formatting fix for mkdocs

* feat: add express-mongo-sanitize
feat: add login/registration rate limiting

* chore: remove unnecessary console log

* wip: remove token handling from localStorage to encrypted DB solution

* refactor: minor change to UserService

* fix mongo query and add keys route to server

* fix backend controllers and simplify schema/crud

* refactor: rename token to key to separate from access/refresh tokens, setTokenDialog -> setKeyDialog

* refactor(schemas): TEndpointOption token -> key

* refactor(api): use new encrypted key retrieval system

* fix(SetKeyDialog): fix key prop error

* fix(abortMiddleware): pass random UUID if messageId is not generated yet for proper error display on frontend

* fix(getUserKey): wrong prop passed in arg, adds error handling

* fix: prevent message without conversationId from saving to DB, prevents branching on the frontend to a new top-level branch

* refactor: change wording of multiple display messages

* refactor(checkExpiry -> checkUserKeyExpiry): move to UserService file

* fix: type imports from common

* refactor(SubmitButton): convert to TS

* refactor(key.ts): change localStorage map key name

* refactor: add new custom tailwind classes to better match openAI colors

* chore: remove unnecessary warning and catch ScreenShot error

* refactor: move userKey frontend logic to hooks and remove use of localStorage and instead query the DB

* refactor: invalidate correct query key, memoize userKey hook, conditionally render SetKeyDialog to avoid unnecessary calls, refactor SubmitButton props and useEffect for showing 'provide key first'

* fix(SetKeyDialog): use enum-like object for expiry values
feat(Dropdown): add optionsClassName to dynamically change dropdown options container classes

* fix: handle edge case where user had provided a key but the server changes to env variable for keys

* refactor(OpenAI/titleConvo): move titling to client to retain authorized credentials in message lifecycle for titling

* fix(azure): handle user_provided keys correctly for azure

* feat: send user Id to OpenAI to differentiate users in completion requests

* refactor(OpenAI/titleConvo): adding tokens helps minimize LLM from using the language in title response

* feat: add delete endpoint for keys

* chore: remove throttling of title

* feat: add 'Data controls' to Settings, add 'Revoke' keys feature in Key Dialog and Data controls

* refactor: reorganize PluginsClient files in langchain format

* feat: use langchain for titling convos

* chore: cleanup titling convo, with fallback to original method, escape braces, use only snippet for language detection

* refactor: move helper functions to appropriate langchain folders for reusability

* fix: userProvidesKey handling for gptPlugins

* fix: frontend handling of plugins key

* chore: cleanup logging and ts-ignore SSE

* fix: forwardRef misuse in DangerButton

* fix(GoogleConfig/FileUpload): localize errors and simplify validation with zod

* fix: cleanup google logging and fix user provided key handling

* chore: remove titling from google

* chore: removing logging from browser endpoint

* wip: fix menu flicker

* feat: useLocalStorage hook

* feat: add Tooltip for UI

* refactor(EndpointMenu): utilize Tooltip and useLocalStorage, remove old 'New Chat' slide-over

* fix(e2e): use testId for endpoint menu trigger

* chore: final touches to EndpointMenu before future refactor to declutter component

* refactor(localization): change select endpoint to open menu and add translations

* chore: add final prop to error message response

* ci: minor edits to facilitate testing

* ci: new e2e test which tests for new key setting/revoking features

packages/data-provider/src/api-endpoints.ts

@@ -1,95 +1,59 @@
-export const user = () => {
-  return '/api/user';
-};
+export const user = () => '/api/user';
 
-export const userPlugins = () => {
-  return '/api/user/plugins';
-};
+export const userPlugins = () => '/api/user/plugins';
 
-export const messages = (conversationId: string, messageId?: string) => {
-  return `/api/messages/${conversationId}${messageId ? `/${messageId}` : ''}`;
-};
+export const messages = (conversationId: string, messageId?: string) =>
+  `/api/messages/${conversationId}${messageId ? `/${messageId}` : ''}`;
 
-export const abortRequest = (endpoint: string) => {
-  return `/api/ask/${endpoint}/abort`;
-};
+const keysEndpoint = '/api/keys';
 
-export const conversations = (pageNumber: string) => {
-  return `/api/convos?pageNumber=${pageNumber}`;
-};
+export const keys = () => keysEndpoint;
 
-export const conversationById = (id: string) => {
-  return `/api/convos/${id}`;
-};
+export const userKeyQuery = (name: string) => `${keysEndpoint}?name=${name}`;
 
-export const updateConversation = () => {
-  return '/api/convos/update';
-};
+export const revokeUserKey = (name: string) => `${keysEndpoint}/${name}`;
 
-export const deleteConversation = () => {
-  return '/api/convos/clear';
-};
+export const revokeAllUserKeys = () => `${keysEndpoint}?all=true`;
 
-export const search = (q: string, pageNumber: string) => {
-  return `/api/search?q=${q}&pageNumber=${pageNumber}`;
-};
+export const abortRequest = (endpoint: string) => `/api/ask/${endpoint}/abort`;
 
-export const searchEnabled = () => {
-  return '/api/search/enable';
-};
+export const conversations = (pageNumber: string) => `/api/convos?pageNumber=${pageNumber}`;
 
-export const presets = () => {
-  return '/api/presets';
-};
+export const conversationById = (id: string) => `/api/convos/${id}`;
 
-export const deletePreset = () => {
-  return '/api/presets/delete';
-};
+export const updateConversation = () => '/api/convos/update';
 
-export const aiEndpoints = () => {
-  return '/api/endpoints';
-};
+export const deleteConversation = () => '/api/convos/clear';
 
-export const tokenizer = () => {
-  return '/api/tokenizer';
-};
+export const search = (q: string, pageNumber: string) =>
+  `/api/search?q=${q}&pageNumber=${pageNumber}`;
 
-export const login = () => {
-  return '/api/auth/login';
-};
+export const searchEnabled = () => '/api/search/enable';
 
-export const logout = () => {
-  return '/api/auth/logout';
-};
+export const presets = () => '/api/presets';
 
-export const register = () => {
-  return '/api/auth/register';
-};
+export const deletePreset = () => '/api/presets/delete';
 
-export const loginFacebook = () => {
-  return '/api/auth/facebook';
-};
+export const aiEndpoints = () => '/api/endpoints';
 
-export const loginGoogle = () => {
-  return '/api/auth/google';
-};
+export const tokenizer = () => '/api/tokenizer';
 
-export const refreshToken = () => {
-  return '/api/auth/refresh';
-};
+export const login = () => '/api/auth/login';
 
-export const requestPasswordReset = () => {
-  return '/api/auth/requestPasswordReset';
-};
+export const logout = () => '/api/auth/logout';
 
-export const resetPassword = () => {
-  return '/api/auth/resetPassword';
-};
+export const register = () => '/api/auth/register';
 
-export const plugins = () => {
-  return '/api/plugins';
-};
+export const loginFacebook = () => '/api/auth/facebook';
 
-export const config = () => {
-  return '/api/config';
-};
+export const loginGoogle = () => '/api/auth/google';
+
+export const refreshToken = () => '/api/auth/refresh';
+
+export const requestPasswordReset = () => '/api/auth/requestPasswordReset';
+
+export const resetPassword = () => '/api/auth/resetPassword';
+
+export const plugins = () => '/api/plugins';
+
+export const config = () => '/api/config';

packages/data-provider/src/data-service.ts

@@ -24,6 +24,14 @@ export function clearAllConversations(): Promise<unknown> {
   return request.post(endpoints.deleteConversation(), { arg: {} });
 }
 
+export function revokeUserKey(name: string): Promise<unknown> {
+  return request.delete(endpoints.revokeUserKey(name));
+}
+
+export function revokeAllUserKeys(): Promise<unknown> {
+  return request.delete(endpoints.revokeAllUserKeys());
+}
+
 export function getMessagesByConvoId(conversationId: string): Promise<s.TMessage[]> {
   return request.get(endpoints.messages(conversationId));
 }
@@ -47,6 +55,15 @@ export function updateMessage(payload: t.TUpdateMessageRequest): Promise<unknown
   return request.put(endpoints.messages(conversationId, messageId), { text });
 }
 
+export function updateUserKey(payload: t.TUpdateUserKeyRequest) {
+  const { value } = payload;
+  if (!value) {
+    throw new Error('value is required');
+  }
+
+  return request.put(endpoints.keys(), payload);
+}
+
 export function getPresets(): Promise<s.TPreset[]> {
   return request.get(endpoints.presets());
 }
@@ -98,9 +115,10 @@ export const register = (payload: t.TRegisterUser) => {
   return request.post(endpoints.register(), payload);
 };
 
-export const refreshToken = () => {
-  return request.post(endpoints.refreshToken());
-};
+export const refreshToken = () => request.post(endpoints.refreshToken());
+
+export const userKeyQuery = (name: string): Promise<t.TCheckUserKeyResponse> =>
+  request.get(endpoints.userKeyQuery(name));
 
 export const getLoginGoogle = () => {
   return request.get(endpoints.loginGoogle());

packages/data-provider/src/react-query-service.ts

@@ -16,6 +16,7 @@ export enum QueryKeys {
   conversation = 'conversation',
   searchEnabled = 'searchEnabled',
   user = 'user',
+  name = 'name', // user key name
   endpoints = 'endpoints',
   presets = 'presets',
   searchResults = 'searchResults',
@@ -121,6 +122,20 @@ export const useUpdateMessageMutation = (
   });
 };
 
+export const useUpdateUserKeysMutation = (): UseMutationResult<
+  t.TUser,
+  unknown,
+  t.TUpdateUserKeyRequest,
+  unknown
+> => {
+  const queryClient = useQueryClient();
+  return useMutation((payload: t.TUpdateUserKeyRequest) => dataService.updateUserKey(payload), {
+    onSuccess: () => {
+      queryClient.invalidateQueries([QueryKeys.name]);
+    },
+  });
+};
+
 export const useDeleteConversationMutation = (
   id?: string,
 ): UseMutationResult<
@@ -150,6 +165,24 @@ export const useClearConversationsMutation = (): UseMutationResult<unknown> => {
   });
 };
 
+export const useRevokeUserKeyMutation = (name: string): UseMutationResult<unknown> => {
+  const queryClient = useQueryClient();
+  return useMutation(() => dataService.revokeUserKey(name), {
+    onSuccess: () => {
+      queryClient.invalidateQueries([QueryKeys.name]);
+    },
+  });
+};
+
+export const useRevokeAllUserKeysMutation = (): UseMutationResult<unknown> => {
+  const queryClient = useQueryClient();
+  return useMutation(() => dataService.revokeAllUserKeys(), {
+    onSuccess: () => {
+      queryClient.invalidateQueries([QueryKeys.name]);
+    },
+  });
+};
+
 export const useGetConversationsQuery = (
   pageNumber: string,
   config?: UseQueryOptions<t.TGetConversationsResponse>,
@@ -315,6 +348,28 @@ export const useRefreshTokenMutation = (): UseMutationResult<
   return useMutation(() => dataService.refreshToken(), {});
 };
 
+export const useUserKeyQuery = (
+  name: string,
+  config?: UseQueryOptions<t.TCheckUserKeyResponse>,
+): QueryObserverResult<t.TCheckUserKeyResponse> => {
+  return useQuery<t.TCheckUserKeyResponse>(
+    [QueryKeys.name, name],
+    () => {
+      if (!name) {
+        return Promise.resolve({ expiresAt: '' });
+      }
+      return dataService.userKeyQuery(name);
+    },
+    {
+      refetchOnWindowFocus: false,
+      refetchOnReconnect: false,
+      refetchOnMount: false,
+      retry: false,
+      ...config,
+    },
+  );
+};
+
 export const useRequestPasswordResetMutation = (): UseMutationResult<
   t.TRequestPasswordResetResponse,
   unknown,

packages/data-provider/src/schemas.ts

@@ -399,7 +399,7 @@ export type TEndpointOption = {
   chatGptLabel?: string | null;
   modelLabel?: string | null;
   jailbreak?: boolean;
-  token?: string | null;
+  key?: string | null;
 };
 
 export const getResponseSender = (endpointOption: TEndpointOption): string => {

packages/data-provider/src/types.ts

@@ -1,4 +1,7 @@
 import type { TResPlugin, TMessage, TConversation, TEndpointOption } from './schemas';
+import type { UseMutationResult } from '@tanstack/react-query';
+
+export type TMutation = UseMutationResult<unknown>;
 
 export * from './schemas';
 
@@ -69,6 +72,12 @@ export type TUpdateMessageRequest = {
   text: string;
 };
 
+export type TUpdateUserKeyRequest = {
+  name: string;
+  value: string;
+  expiresAt: string;
+};
+
 export type TUpdateConversationRequest = {
   conversationId: string;
   title: string;
@@ -177,6 +186,10 @@ export type TRefreshTokenResponse = {
   user: TUser;
 };
 
+export type TCheckUserKeyResponse = {
+  expiresAt: string;
+};
+
 export type TRequestPasswordResetResponse = {
   link?: string;
   message?: string;