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

client/src/components/Input/SetKeyDialog/GoogleConfig.tsx

@@ -0,0 +1,35 @@
+import React from 'react';
+import { object, string } from 'zod';
+import type { TConfigProps } from '~/common';
+import FileUpload from '../EndpointMenu/FileUpload';
+import { useLocalize } from '~/hooks';
+
+const CredentialsSchema = object({
+  client_email: string().email().min(3),
+  project_id: string().min(3),
+  private_key: string().min(601),
+});
+
+const validateCredentials = (credentials: Record<string, unknown>) => {
+  const result = CredentialsSchema.safeParse(credentials);
+  return result.success;
+};
+
+const GoogleConfig = ({ setUserKey }: Pick<TConfigProps, 'setUserKey'>) => {
+  const localize = useLocalize();
+  return (
+    <FileUpload
+      id="googleKey"
+      className="w-full"
+      text={localize('com_endpoint_config_key_import_json_key')}
+      successText={localize('com_endpoint_config_key_import_json_key_success')}
+      invalidText={localize('com_endpoint_config_key_import_json_key_invalid')}
+      validator={validateCredentials}
+      onFileSelected={(data) => {
+        setUserKey(JSON.stringify(data));
+      }}
+    />
+  );
+};
+
+export default GoogleConfig;

client/src/components/Input/SetKeyDialog/HelpText.tsx

@@ -0,0 +1,85 @@
+import React from 'react';
+import { useLocalize } from '~/hooks';
+
+function HelpText({ endpoint }: { endpoint: string }) {
+  const localize = useLocalize();
+  const textMap = {
+    bingAI: (
+      <small className="break-all text-gray-600">
+        {localize('com_endpoint_config_key_get_edge_key')}{' '}
+        <a
+          target="_blank"
+          href="https://www.bing.com"
+          rel="noreferrer"
+          className="text-blue-600 underline"
+        >
+          https://www.bing.com
+        </a>
+        {'. '}
+        {localize('com_endpoint_config_key_get_edge_key_dev_tool')}{' '}
+        <a
+          target="_blank"
+          href="https://github.com/waylaidwanderer/node-chatgpt-api/issues/378#issuecomment-1559868368"
+          rel="noreferrer"
+          className="text-blue-600 underline"
+        >
+          {localize('com_endpoint_config_key_edge_instructions')}
+        </a>{' '}
+        {localize('com_endpoint_config_key_edge_full_token_string')}
+      </small>
+    ),
+    chatGPTBrowser: (
+      <small className="break-all text-gray-600">
+        {localize('com_endpoint_config_key_chatgpt')}{' '}
+        <a
+          target="_blank"
+          href="https://chat.openai.com"
+          rel="noreferrer"
+          className="text-blue-600 underline"
+        >
+          https://chat.openai.com
+        </a>
+        {', '}
+        {localize('com_endpoint_config_key_chatgpt_then_visit')}{' '}
+        <a
+          target="_blank"
+          href="https://chat.openai.com/api/auth/session"
+          rel="noreferrer"
+          className="text-blue-600 underline"
+        >
+          https://chat.openai.com/api/auth/session
+        </a>
+        {'. '}
+        {localize('com_endpoint_config_key_chatgpt_copy_token')}
+      </small>
+    ),
+    google: (
+      <small className="break-all text-gray-600">
+        {localize('com_endpoint_config_key_google_need_to')}{' '}
+        <a
+          target="_blank"
+          href="https://console.cloud.google.com/vertex-ai"
+          rel="noreferrer"
+          className="text-blue-600 underline"
+        >
+          {localize('com_endpoint_config_key_google_vertex_ai')}
+        </a>{' '}
+        {localize('com_endpoint_config_key_google_vertex_api')}{' '}
+        <a
+          target="_blank"
+          href="https://console.cloud.google.com/projectselector/iam-admin/serviceaccounts/create?walkthrough_id=iam--create-service-account#step_index=1"
+          rel="noreferrer"
+          className="text-blue-600 underline"
+        >
+          {localize('com_endpoint_config_key_google_service_account')}
+        </a>
+        {'. '}
+        {localize('com_endpoint_config_key_google_vertex_api_role')}
+      </small>
+    ),
+  };
+
+  return textMap[endpoint] || null;
+}
+
+export default React.memo(HelpText);

client/src/components/Input/SetKeyDialog/InputWithLabel.tsx

@@ -0,0 +1,38 @@
+import React, { ChangeEvent, FC } from 'react';
+import { Input, Label } from '~/components';
+import { cn, defaultTextPropsLabel, removeFocusOutlines } from '~/utils/';
+import { useLocalize } from '~/hooks';
+
+interface InputWithLabelProps {
+  value: string;
+  onChange: (event: ChangeEvent<HTMLInputElement>) => void;
+  label: string;
+  id: string;
+}
+
+const InputWithLabel: FC<InputWithLabelProps> = ({ value, onChange, label, id }) => {
+  const localize = useLocalize();
+  return (
+    <>
+      <Label htmlFor={id} className="text-left text-sm font-medium">
+        {label}
+        <br />
+      </Label>
+
+      <Input
+        id={id}
+        data-testid={`input-${id}`}
+        value={value ?? ''}
+        onChange={onChange}
+        placeholder={`${localize('com_endpoint_config_value')} ${label}`}
+        className={cn(
+          defaultTextPropsLabel,
+          'flex h-10 max-h-10 w-full resize-none px-3 py-2',
+          removeFocusOutlines,
+        )}
+      />
+    </>
+  );
+};
+
+export default InputWithLabel;

client/src/components/Input/SetKeyDialog/OpenAIConfig.tsx

@@ -0,0 +1,127 @@
+/* eslint-disable react-hooks/exhaustive-deps */
+import React, { useEffect, useState } from 'react';
+// TODO: Temporarily remove checkbox until Plugins solution for Azure is figured out
+// import * as Checkbox from '@radix-ui/react-checkbox';
+// import { CheckIcon } from '@radix-ui/react-icons';
+import InputWithLabel from './InputWithLabel';
+import type { TConfigProps } from '~/common';
+
+function isJson(str: string) {
+  try {
+    JSON.parse(str);
+  } catch (e) {
+    return false;
+  }
+  return true;
+}
+
+const OpenAIConfig = ({ userKey, setUserKey, endpoint }: TConfigProps) => {
+  const [showPanel, setShowPanel] = useState(endpoint === 'azureOpenAI');
+
+  useEffect(() => {
+    if (isJson(userKey)) {
+      setShowPanel(true);
+    }
+    setUserKey('');
+  }, []);
+
+  useEffect(() => {
+    if (!showPanel && isJson(userKey)) {
+      setUserKey('');
+    }
+  }, [showPanel]);
+
+  function getAzure(name: string) {
+    if (isJson(userKey)) {
+      const newKey = JSON.parse(userKey);
+      return newKey[name];
+    } else {
+      return '';
+    }
+  }
+
+  function setAzure(name: string, value: number | string | boolean) {
+    let newKey = {};
+    if (isJson(userKey)) {
+      newKey = JSON.parse(userKey);
+    }
+    newKey[name] = value;
+
+    setUserKey(JSON.stringify(newKey));
+  }
+  return (
+    <>
+      {!showPanel ? (
+        <>
+          <InputWithLabel
+            id={endpoint}
+            value={userKey ?? ''}
+            onChange={(e: { target: { value: string } }) => setUserKey(e.target.value ?? '')}
+            label={'OpenAI API Key'}
+          />
+        </>
+      ) : (
+        <>
+          <InputWithLabel
+            id={'instanceNameLabel'}
+            value={getAzure('azureOpenAIApiInstanceName') ?? ''}
+            onChange={(e: { target: { value: string } }) =>
+              setAzure('azureOpenAIApiInstanceName', e.target.value ?? '')
+            }
+            label={'Azure OpenAI Instance Name'}
+          />
+
+          <InputWithLabel
+            id={'deploymentNameLabel'}
+            value={getAzure('azureOpenAIApiDeploymentName') ?? ''}
+            onChange={(e: { target: { value: string } }) =>
+              setAzure('azureOpenAIApiDeploymentName', e.target.value ?? '')
+            }
+            label={'Azure OpenAI Deployment Name'}
+          />
+
+          <InputWithLabel
+            id={'versionLabel'}
+            value={getAzure('azureOpenAIApiVersion') ?? ''}
+            onChange={(e: { target: { value: string } }) =>
+              setAzure('azureOpenAIApiVersion', e.target.value ?? '')
+            }
+            label={'Azure OpenAI API Version'}
+          />
+
+          <InputWithLabel
+            id={'apiKeyLabel'}
+            value={getAzure('azureOpenAIApiKey') ?? ''}
+            onChange={(e: { target: { value: string } }) =>
+              setAzure('azureOpenAIApiKey', e.target.value ?? '')
+            }
+            label={'Azure OpenAI API Key'}
+          />
+        </>
+      )}
+      {/* { endpoint === 'gptPlugins' && (
+        <div className="flex items-center">
+          <Checkbox.Root
+            className="flex h-[20px] w-[20px] appearance-none items-center justify-center rounded-[4px] bg-gray-100 text-white outline-none hover:bg-gray-200 dark:bg-gray-700 dark:hover:bg-gray-900"
+            id="azureOpenAI"
+            checked={showPanel}
+            onCheckedChange={() => setShowPanel(!showPanel)}
+          >
+            <Checkbox.Indicator className="flex h-[20px] w-[20px] items-center justify-center rounded-[3.5px] bg-green-600">
+              <CheckIcon />
+            </Checkbox.Indicator>
+          </Checkbox.Root>
+
+          <label
+            className="pl-[8px] text-[15px] leading-none dark:text-white"
+            htmlFor="azureOpenAI"
+          >
+          Use Azure OpenAI.
+          </label>
+        </div>
+      )} */}
+    </>
+  );
+};
+
+export default OpenAIConfig;

client/src/components/Input/SetKeyDialog/OtherConfig.tsx

@@ -0,0 +1,18 @@
+import React from 'react';
+import InputWithLabel from './InputWithLabel';
+import type { TConfigProps } from '~/common';
+import { useLocalize } from '~/hooks';
+
+const OtherConfig = ({ userKey, setUserKey, endpoint }: TConfigProps) => {
+  const localize = useLocalize();
+  return (
+    <InputWithLabel
+      id={endpoint}
+      value={userKey ?? ''}
+      onChange={(e: React.ChangeEvent<HTMLInputElement>) => setUserKey(e.target.value ?? '')}
+      label={localize('com_endpoint_config_key_name')}
+    />
+  );
+};
+
+export default OtherConfig;

client/src/components/Input/SetKeyDialog/SetKeyDialog.tsx

@@ -0,0 +1,103 @@
+import React, { useState } from 'react';
+import type { TDialogProps } from '~/common';
+import { Dialog, Dropdown } from '~/components/ui';
+import DialogTemplate from '~/components/ui/DialogTemplate';
+import { RevokeKeysButton } from '~/components/Nav';
+import { cn, defaultTextProps, removeFocusOutlines, alternateName } from '~/utils';
+import { useUserKey, useLocalize } from '~/hooks';
+import GoogleConfig from './GoogleConfig';
+import OpenAIConfig from './OpenAIConfig';
+import OtherConfig from './OtherConfig';
+import HelpText from './HelpText';
+
+const endpointComponents = {
+  google: GoogleConfig,
+  openAI: OpenAIConfig,
+  azureOpenAI: OpenAIConfig,
+  gptPlugins: OpenAIConfig,
+  default: OtherConfig,
+};
+
+const EXPIRY = {
+  THIRTY_MINUTES: { display: 'in 30 minutes', value: 30 * 60 * 1000 },
+  TWO_HOURS: { display: 'in 2 hours', value: 2 * 60 * 60 * 1000 },
+  TWELVE_HOURS: { display: 'in 12 hours', value: 12 * 60 * 60 * 1000 },
+  ONE_DAY: { display: 'in 1 day', value: 24 * 60 * 60 * 1000 },
+  ONE_WEEK: { display: 'in 7 days', value: 7 * 24 * 60 * 60 * 1000 },
+  ONE_MONTH: { display: 'in 30 days', value: 30 * 24 * 60 * 60 * 1000 },
+};
+
+const SetKeyDialog = ({
+  open,
+  onOpenChange,
+  endpoint,
+}: Pick<TDialogProps, 'open' | 'onOpenChange'> & {
+  endpoint: string;
+}) => {
+  const [userKey, setUserKey] = useState('');
+  const [expiresAtLabel, setExpiresAtLabel] = useState(EXPIRY.TWELVE_HOURS.display);
+  const { getExpiry, saveUserKey } = useUserKey(endpoint);
+  const localize = useLocalize();
+
+  const expirationOptions = Object.values(EXPIRY);
+
+  const handleExpirationChange = (label: string) => {
+    setExpiresAtLabel(label);
+  };
+
+  const submit = () => {
+    const selectedOption = expirationOptions.find((option) => option.display === expiresAtLabel);
+    const expiresAt = Date.now() + (selectedOption ? selectedOption.value : 0);
+    saveUserKey(userKey, expiresAt);
+    onOpenChange(false);
+    setUserKey('');
+  };
+
+  const EndpointComponent = endpointComponents[endpoint] ?? endpointComponents['default'];
+  const expiryTime = getExpiry();
+
+  return (
+    <Dialog open={open} onOpenChange={onOpenChange}>
+      <DialogTemplate
+        title={`${localize('com_endpoint_config_key_for')} ${alternateName[endpoint] ?? endpoint}`}
+        className="w-full max-w-[650px] sm:w-3/4 md:w-3/4 lg:w-3/4"
+        main={
+          <div className="grid w-full items-center gap-2">
+            <small className="text-red-600">
+              {`${localize('com_endpoint_config_key_encryption')} ${
+                !expiryTime
+                  ? localize('com_endpoint_config_key_expiry')
+                  : `${new Date(expiryTime).toLocaleString()}`
+              }`}
+            </small>
+            <Dropdown
+              label="Expires "
+              value={expiresAtLabel}
+              onChange={handleExpirationChange}
+              options={expirationOptions.map((option) => option.display)}
+              className={cn(
+                defaultTextProps,
+                'flex h-full w-full resize-none',
+                removeFocusOutlines,
+              )}
+              optionsClassName="max-h-72"
+              containerClassName="flex w-1/2 md:w-1/3 resize-none z-[51]"
+            />
+            <EndpointComponent userKey={userKey} setUserKey={setUserKey} endpoint={endpoint} />
+            <HelpText endpoint={endpoint} />
+          </div>
+        }
+        selection={{
+          selectHandler: submit,
+          selectClasses: 'bg-green-600 hover:bg-green-700 dark:hover:bg-green-800 text-white',
+          selectText: localize('com_ui_submit'),
+        }}
+        leftButtons={
+          <RevokeKeysButton endpoint={endpoint} showText={false} disabled={!expiryTime} />
+        }
+      />
+    </Dialog>
+  );
+};
+
+export default SetKeyDialog;

client/src/components/Input/SetKeyDialog/index.ts

@@ -0,0 +1 @@
+export { default as SetKeyDialog } from './SetKeyDialog';