diff --git a/public/scripts/extensions/infinity-context/index.js b/public/scripts/extensions/infinity-context/index.js
index eae7ad4ca..97af2dec6 100644
--- a/public/scripts/extensions/infinity-context/index.js
+++ b/public/scripts/extensions/infinity-context/index.js
@@ -1,7 +1,7 @@
 import { saveSettingsDebounced, getCurrentChatId } from "../../../script.js";
 import { getApiUrl, extension_settings } from "../../extensions.js";
 import { splitRecursive } from "../../utils.js";
-export { MODULE_NAME, chromadb_interceptGeneration };
+export { MODULE_NAME };
 
 const MODULE_NAME = 'chromadb';
 
@@ -125,7 +125,7 @@ window.chromadb_interceptGeneration = async (chat) => {
             await addMessages(currentChatId, messagesToStore);
         }
 
-        const lastMessage = messagesToKeep[messagesToKeep.length - 1];
+        const lastMessage = chat[chat.length - 1];
 
         if (lastMessage) {
             const queriedMessages = await queryMessages(currentChatId, lastMessage.mes);
diff --git a/public/scripts/utils.js b/public/scripts/utils.js
index 158f95471..915c60096 100644
--- a/public/scripts/utils.js
+++ b/public/scripts/utils.js
@@ -262,4 +262,4 @@ export function splitRecursive(input, length, delimitiers = ['\n\n', '\n', ' ',
         result.push(currentChunk);
     }
     return result;
-}
\ No newline at end of file
+}