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

api/server/routes/ask/anthropic.js

@@ -87,7 +87,7 @@ router.post(
         getAbortData,
       );
 
-      const { client } = initializeClient(req, endpointOption);
+      const { client } = await initializeClient(req, endpointOption);
 
       let response = await client.sendMessage(text, {
         getIds,
@@ -135,7 +135,7 @@ router.post(
         conversationId,
         sender: getResponseSender(endpointOption),
         messageId: responseMessageId,
-        parentMessageId: userMessageId,
+        parentMessageId: userMessageId ?? parentMessageId,
       });
     }
   },

api/server/routes/ask/askChatGPTBrowser.js

@@ -38,7 +38,7 @@ router.post('/', requireJwtAuth, setHeaders, async (req, res) => {
   // build endpoint option
   const endpointOption = {
     model: req.body?.model ?? 'text-davinci-002-render-sha',
-    token: req.body?.token ?? null,
+    key: req.body?.key ?? null,
   };
 
   // const availableModels = getChatGPTBrowserModels();

api/server/routes/ask/bingAI.js

@@ -45,7 +45,7 @@ router.post('/', requireJwtAuth, setHeaders, async (req, res) => {
       systemMessage: req.body?.systemMessage ?? null,
       context: req.body?.context ?? null,
       toneStyle: req.body?.toneStyle ?? 'creative',
-      token: req.body?.token ?? null,
+      key: req.body?.key ?? null,
     };
   } else {
     endpointOption = {
@@ -56,7 +56,7 @@ router.post('/', requireJwtAuth, setHeaders, async (req, res) => {
       clientId: req.body?.clientId ?? null,
       invocationId: req.body?.invocationId ?? null,
       toneStyle: req.body?.toneStyle ?? 'creative',
-      token: req.body?.token ?? null,
+      key: req.body?.key ?? null,
     };
   }
 
@@ -139,6 +139,7 @@ const ask = async ({
   try {
     let response = await askBing({
       text,
+      userId: req.user.id,
       parentMessageId: userParentMessageId,
       conversationId: bingConversationId ?? conversationId,
       ...endpointOption,

api/server/routes/ask/google.js

@@ -1,9 +1,10 @@
 const express = require('express');
 const router = express.Router();
 const crypto = require('crypto');
-const { titleConvo, GoogleClient } = require('../../../app');
+const { GoogleClient } = require('../../../app');
 const { saveMessage, getConvoTitle, saveConvo, getConvo } = require('../../../models');
 const { handleError, sendMessage, createOnProgress } = require('../../utils');
+const { getUserKey, checkUserKeyExpiry } = require('../../services/UserService');
 const { requireJwtAuth, setHeaders } = require('../../middleware');
 
 router.post('/', requireJwtAuth, setHeaders, async (req, res) => {
@@ -19,7 +20,7 @@ router.post('/', requireJwtAuth, setHeaders, async (req, res) => {
   const endpointOption = {
     examples: req.body?.examples ?? [{ input: { content: '' }, output: { content: '' } }],
     promptPrefix: req.body?.promptPrefix ?? null,
-    token: req.body?.token ?? null,
+    key: req.body?.key ?? null,
     modelOptions: {
       model: req.body?.model ?? 'chat-bison',
       modelLabel: req.body?.modelLabel ?? null,
@@ -88,17 +89,22 @@ const ask = async ({ text, endpointOption, parentMessageId = null, conversationI
 
     const abortController = new AbortController();
 
+    const isUserProvided = process.env.PALM_KEY === 'user_provided';
+
     let key;
-    if (endpointOption.token) {
-      key = JSON.parse(endpointOption.token);
-      delete endpointOption.token;
+    if (endpointOption.key && isUserProvided) {
+      checkUserKeyExpiry(
+        endpointOption.key,
+        'Your GOOGLE_TOKEN has expired. Please provide your token again.',
+      );
+      key = await getUserKey({ userId: req.user.id, name: 'google' });
+      key = JSON.parse(key);
+      delete endpointOption.key;
       console.log('Using service account key provided by User for PaLM models');
     }
 
     try {
-      if (!key) {
-        key = require('../../../data/auth.json');
-      }
+      key = require('../../../data/auth.json');
     } catch (e) {
       console.log('No \'auth.json\' file (service account key) found in /api/data/ for PaLM models');
     }
@@ -146,14 +152,6 @@ const ask = async ({ text, endpointOption, parentMessageId = null, conversationI
       responseMessage: response,
     });
     res.end();
-
-    if (parentMessageId == '00000000-0000-0000-0000-000000000000') {
-      const title = await titleConvo({ text, response });
-      await saveConvo(req.user.id, {
-        conversationId,
-        title,
-      });
-    }
   } catch (error) {
     console.error(error);
     const errorMessage = {

api/server/routes/ask/gptPlugins.js

@@ -158,7 +158,7 @@ router.post(
 
     try {
       endpointOption.tools = await validateTools(user, endpointOption.tools);
-      const { client, azure, openAIApiKey } = initializeClient(req, endpointOption);
+      const { client } = await initializeClient(req, endpointOption);
 
       let response = await client.sendMessage(text, {
         user,
@@ -204,14 +204,14 @@ router.post(
         responseMessage: response,
       });
       res.end();
-      addTitle(req, {
-        text,
-        newConvo,
-        response,
-        openAIApiKey,
-        parentMessageId,
-        azure: !!azure,
-      });
+
+      if (parentMessageId == '00000000-0000-0000-0000-000000000000' && newConvo) {
+        addTitle(req, {
+          text,
+          response,
+          client,
+        });
+      }
     } catch (error) {
       const partialText = getPartialText();
       handleAbortError(res, req, error, {
@@ -219,7 +219,7 @@ router.post(
         conversationId,
         sender: getResponseSender(endpointOption),
         messageId: responseMessageId,
-        parentMessageId: userMessageId,
+        parentMessageId: userMessageId ?? parentMessageId,
       });
     }
   },

api/server/routes/ask/openAI.js

@@ -94,7 +94,7 @@ router.post(
     );
 
     try {
-      const { client, openAIApiKey } = initializeClient(req, endpointOption);
+      const { client } = await initializeClient(req, endpointOption);
 
       let response = await client.sendMessage(text, {
         user,
@@ -136,14 +136,13 @@ router.post(
       });
       res.end();
 
-      addTitle(req, {
-        text,
-        newConvo,
-        response,
-        openAIApiKey,
-        parentMessageId,
-        azure: endpointOption.endpoint === 'azureOpenAI',
-      });
+      if (parentMessageId == '00000000-0000-0000-0000-000000000000' && newConvo) {
+        addTitle(req, {
+          text,
+          response,
+          client,
+        });
+      }
     } catch (error) {
       const partialText = getPartialText();
       handleAbortError(res, req, error, {
@@ -151,7 +150,7 @@ router.post(
         conversationId,
         sender: getResponseSender(endpointOption),
         messageId: responseMessageId,
-        parentMessageId: userMessageId,
+        parentMessageId: userMessageId ?? parentMessageId,
       });
     }
   },