🛠️ fix: Minor Fixes in Message, Ask/EditController, `OpenAIClient…

…`, and `countTokens` (danny-avila#1463) * fix(Message): avoid overwriting unprovided properties * fix(OpenAIClient): return intermediateReply on user abort * fix(AskController): do not send/save final message if abort was triggered * fix(countTokens): avoid fetching remote registry and exclusively use cl100k_base or p50k_base weights for token counting * refactor(Message/messageSchema): rely on messageSchema for default values when saving messages * fix(EditController): do not send/save final message if abort was triggered * fix(config/helpers): fix module resolution error
aitok-ai · Dec 30, 2023 · 98017a8 · 98017a8
1 parent 7a2ad12
commit 98017a8
Show file tree

Hide file tree

Showing 6 changed files with 35 additions and 30 deletions.
diff --git a/app/clients/OpenAIClient.js b/app/clients/OpenAIClient.js
@@ -847,7 +847,7 @@ ${convo}
         err?.message?.includes('abort') ||
         (err instanceof OpenAI.APIError && err?.message?.includes('abort'))
       ) {
-        return '';
+        return intermediateReply;
       }
       if (
         err?.message?.includes(

diff --git a/models/Message.js b/models/Message.js
@@ -15,16 +15,16 @@ module.exports = {
     parentMessageId,
     sender,
     text,
-    isCreatedByUser = false,
+    isCreatedByUser,
     error,
     unfinished,
     files,
-    isEdited = false,
-    finish_reason = null,
-    tokenCount = null,
-    plugin = null,
-    plugins = null,
-    model = null,
+    isEdited,
+    finish_reason,
+    tokenCount,
+    plugin,
+    plugins,
+    model,
   }) {
     try {
       const validConvoId = idSchema.safeParse(conversationId);

diff --git a/models/schema/messageSchema.js b/models/schema/messageSchema.js
@@ -21,6 +21,7 @@ const messageSchema = mongoose.Schema(
     },
     model: {
       type: String,
+      default: null,
     },
     conversationSignature: {
       type: String,

diff --git a/server/controllers/AskController.js b/server/controllers/AskController.js
@@ -118,16 +118,19 @@ const AskController = async (req, res, next, initializeClient, addTitle) => {
       delete userMessage.image_urls;
     }
 
-    sendMessage(res, {
-      title: await getConvoTitle(user, conversationId),
-      final: true,
-      conversation: await getConvo(user, conversationId),
-      requestMessage: userMessage,
-      responseMessage: response,
-    });
-    res.end();
+    if (!abortController.signal.aborted) {
+      sendMessage(res, {
+        title: await getConvoTitle(user, conversationId),
+        final: true,
+        conversation: await getConvo(user, conversationId),
+        requestMessage: userMessage,
+        responseMessage: response,
+      });
+      res.end();
+
+      await saveMessage({ ...response, user });
+    }
 
-    await saveMessage({ ...response, user });
     await saveMessage(userMessage);
 
     if (addTitle && parentMessageId === '00000000-0000-0000-0000-000000000000' && newConvo) {

diff --git a/server/controllers/EditController.js b/server/controllers/EditController.js
@@ -112,16 +112,18 @@ const EditController = async (req, res, next, initializeClient) => {
       response = { ...response, ...metadata };
     }
 
-    await saveMessage({ ...response, user });
+    if (!abortController.signal.aborted) {
+      sendMessage(res, {
+        title: await getConvoTitle(user, conversationId),
+        final: true,
+        conversation: await getConvo(user, conversationId),
+        requestMessage: userMessage,
+        responseMessage: response,
+      });
+      res.end();
 
-    sendMessage(res, {
-      title: await getConvoTitle(user, conversationId),
-      final: true,
-      conversation: await getConvo(user, conversationId),
-      requestMessage: userMessage,
-      responseMessage: response,
-    });
-    res.end();
+      await saveMessage({ ...response, user });
+    }
   } catch (error) {
     const partialText = getPartialText();
     handleAbortError(res, req, error, {

diff --git a/server/utils/countTokens.js b/server/utils/countTokens.js
@@ -1,13 +1,12 @@
-const { load } = require('tiktoken/load');
 const { Tiktoken } = require('tiktoken/lite');
-const registry = require('tiktoken/registry.json');
-const models = require('tiktoken/model_to_encoding.json');
+const p50k_base = require('tiktoken/encoders/p50k_base.json');
+const cl100k_base = require('tiktoken/encoders/cl100k_base.json');
 const logger = require('~/config/winston');
 
 const countTokens = async (text = '', modelName = 'gpt-3.5-turbo') => {
   let encoder = null;
   try {
-    const model = await load(registry[models[modelName]]);
+    const model = modelName.includes('text-davinci-003') ? p50k_base : cl100k_base;
     encoder = new Tiktoken(model.bpe_ranks, model.special_tokens, model.pat_str);
     const tokens = encoder.encode(text);
     encoder.free();