diff --git a/.github/workflows/dev-build.yaml b/.github/workflows/dev-build.yaml
index bf9a1e67fc3267dd314457f32973cb253db6e9eb..bcd509b5c853398c36ea13e3e29d9aa640a1840f 100644
--- a/.github/workflows/dev-build.yaml
+++ b/.github/workflows/dev-build.yaml
@@ -6,7 +6,7 @@ concurrency:
on:
push:
- branches: ['agent-ui-animations'] # put your current branch to create a build. Core team only.
+ branches: ['3069-tokenizer-collector-improvements'] # put your current branch to create a build. Core team only.
paths-ignore:
- '**.md'
- 'cloud-deployments/*'
diff --git a/collector/processLink/convert/generic.js b/collector/processLink/convert/generic.js
index a5eb20ca945e23a6dda343ec0ad6a48b92da0cbd..4afb9b95483a6c80d2e24f9a5382a92505834613 100644
--- a/collector/processLink/convert/generic.js
+++ b/collector/processLink/convert/generic.js
@@ -41,7 +41,7 @@ async function scrapeGenericUrl(link, textOnly = false) {
published: new Date().toLocaleString(),
wordCount: content.split(" ").length,
pageContent: content,
- token_count_estimate: tokenizeString(content).length,
+ token_count_estimate: tokenizeString(content),
};
const document = writeToServerDocuments(
diff --git a/collector/processRawText/index.js b/collector/processRawText/index.js
index d435c9e7e07efc66cb440e5e856d5832fea771e9..a29eb63c37de09cd0d15a4d0190692124881f069 100644
--- a/collector/processRawText/index.js
+++ b/collector/processRawText/index.js
@@ -55,7 +55,7 @@ async function processRawText(textContent, metadata) {
published: METADATA_KEYS.possible.published(metadata),
wordCount: textContent.split(" ").length,
pageContent: textContent,
- token_count_estimate: tokenizeString(textContent).length,
+ token_count_estimate: tokenizeString(textContent),
};
const document = writeToServerDocuments(
diff --git a/collector/processSingleFile/convert/asAudio.js b/collector/processSingleFile/convert/asAudio.js
index 170426e4068ffadd880895f4b710f8eb653593c4..5f033af74a9b4293fcd84062757d58b822bd73fb 100644
--- a/collector/processSingleFile/convert/asAudio.js
+++ b/collector/processSingleFile/convert/asAudio.js
@@ -56,7 +56,7 @@ async function asAudio({ fullFilePath = "", filename = "", options = {} }) {
published: createdDate(fullFilePath),
wordCount: content.split(" ").length,
pageContent: content,
- token_count_estimate: tokenizeString(content).length,
+ token_count_estimate: tokenizeString(content),
};
const document = writeToServerDocuments(
diff --git a/collector/processSingleFile/convert/asDocx.js b/collector/processSingleFile/convert/asDocx.js
index b0fbd8843ed84aa198f6b35d5f457c1c03d4bcc2..d33a46b9433dd49bb44e58eba84b93f4a7be0c84 100644
--- a/collector/processSingleFile/convert/asDocx.js
+++ b/collector/processSingleFile/convert/asDocx.js
@@ -42,7 +42,7 @@ async function asDocX({ fullFilePath = "", filename = "" }) {
published: createdDate(fullFilePath),
wordCount: content.split(" ").length,
pageContent: content,
- token_count_estimate: tokenizeString(content).length,
+ token_count_estimate: tokenizeString(content),
};
const document = writeToServerDocuments(
diff --git a/collector/processSingleFile/convert/asEPub.js b/collector/processSingleFile/convert/asEPub.js
index 827e3c3af4567aacf417d376158f2bee7cb9479d..51bb20c809b31946cc239c9505291f8d27f14cc2 100644
--- a/collector/processSingleFile/convert/asEPub.js
+++ b/collector/processSingleFile/convert/asEPub.js
@@ -40,7 +40,7 @@ async function asEPub({ fullFilePath = "", filename = "" }) {
published: createdDate(fullFilePath),
wordCount: content.split(" ").length,
pageContent: content,
- token_count_estimate: tokenizeString(content).length,
+ token_count_estimate: tokenizeString(content),
};
const document = writeToServerDocuments(
diff --git a/collector/processSingleFile/convert/asMbox.js b/collector/processSingleFile/convert/asMbox.js
index 4adde23ec93625d4f7767d59e36fe3f43796bdd7..48de60fa37ad93e15e93b40c75414449b0684775 100644
--- a/collector/processSingleFile/convert/asMbox.js
+++ b/collector/processSingleFile/convert/asMbox.js
@@ -53,7 +53,7 @@ async function asMbox({ fullFilePath = "", filename = "" }) {
published: createdDate(fullFilePath),
wordCount: content.split(" ").length,
pageContent: content,
- token_count_estimate: tokenizeString(content).length,
+ token_count_estimate: tokenizeString(content),
};
item++;
diff --git a/collector/processSingleFile/convert/asOfficeMime.js b/collector/processSingleFile/convert/asOfficeMime.js
index b6c3c0601f9ce7fcc494e7ea11407c1ec41b2285..09e320d168908925ccb055dfd0de3a23ec6e0663 100644
--- a/collector/processSingleFile/convert/asOfficeMime.js
+++ b/collector/processSingleFile/convert/asOfficeMime.js
@@ -38,7 +38,7 @@ async function asOfficeMime({ fullFilePath = "", filename = "" }) {
published: createdDate(fullFilePath),
wordCount: content.split(" ").length,
pageContent: content,
- token_count_estimate: tokenizeString(content).length,
+ token_count_estimate: tokenizeString(content),
};
const document = writeToServerDocuments(
diff --git a/collector/processSingleFile/convert/asPDF/index.js b/collector/processSingleFile/convert/asPDF/index.js
index bf14516419e266ade79b245fece87287198b0ed0..e3e42d3bd7039ed9e63b9c7b9705202e6d145bfe 100644
--- a/collector/processSingleFile/convert/asPDF/index.js
+++ b/collector/processSingleFile/convert/asPDF/index.js
@@ -49,7 +49,7 @@ async function asPdf({ fullFilePath = "", filename = "" }) {
published: createdDate(fullFilePath),
wordCount: content.split(" ").length,
pageContent: content,
- token_count_estimate: tokenizeString(content).length,
+ token_count_estimate: tokenizeString(content),
};
const document = writeToServerDocuments(
diff --git a/collector/processSingleFile/convert/asTxt.js b/collector/processSingleFile/convert/asTxt.js
index 53987f247dfeabaef3a89504f65c3e945985cb7e..bc95969e14db807e3f2754b187669863c042fdb0 100644
--- a/collector/processSingleFile/convert/asTxt.js
+++ b/collector/processSingleFile/convert/asTxt.js
@@ -38,7 +38,7 @@ async function asTxt({ fullFilePath = "", filename = "" }) {
published: createdDate(fullFilePath),
wordCount: content.split(" ").length,
pageContent: content,
- token_count_estimate: tokenizeString(content).length,
+ token_count_estimate: tokenizeString(content),
};
const document = writeToServerDocuments(
diff --git a/collector/processSingleFile/convert/asXlsx.js b/collector/processSingleFile/convert/asXlsx.js
index f21c6f1d9bf93779643197515ee5275cfbdb80a9..ca9b8ebac9dbdc8580bff538fe0a1c0e0d7f4773 100644
--- a/collector/processSingleFile/convert/asXlsx.js
+++ b/collector/processSingleFile/convert/asXlsx.js
@@ -67,7 +67,7 @@ async function asXlsx({ fullFilePath = "", filename = "" }) {
published: createdDate(fullFilePath),
wordCount: content.split(/\s+/).length,
pageContent: content,
- token_count_estimate: tokenizeString(content).length,
+ token_count_estimate: tokenizeString(content),
};
const document = writeToServerDocuments(
diff --git a/collector/utils/extensions/Confluence/index.js b/collector/utils/extensions/Confluence/index.js
index c8ab9b03c307afc4e10cb354ee1b24ba396cec02..e0699222a99d9b3a68edfaceea3b3e0e3441b810 100644
--- a/collector/utils/extensions/Confluence/index.js
+++ b/collector/utils/extensions/Confluence/index.js
@@ -104,7 +104,7 @@ async function loadConfluence(
published: new Date().toLocaleString(),
wordCount: doc.pageContent.split(" ").length,
pageContent: doc.pageContent,
- token_count_estimate: tokenizeString(doc.pageContent).length,
+ token_count_estimate: tokenizeString(doc.pageContent),
};
console.log(
diff --git a/collector/utils/extensions/RepoLoader/GithubRepo/index.js b/collector/utils/extensions/RepoLoader/GithubRepo/index.js
index 41147278cdd3d3e8f2fc0a1458901629b9aa2f05..b2efe221146a50aa0c657cd9b713d05e7a76714c 100644
--- a/collector/utils/extensions/RepoLoader/GithubRepo/index.js
+++ b/collector/utils/extensions/RepoLoader/GithubRepo/index.js
@@ -66,7 +66,7 @@ async function loadGithubRepo(args, response) {
published: new Date().toLocaleString(),
wordCount: doc.pageContent.split(" ").length,
pageContent: doc.pageContent,
- token_count_estimate: tokenizeString(doc.pageContent).length,
+ token_count_estimate: tokenizeString(doc.pageContent),
};
console.log(
`[Github Loader]: Saving ${doc.metadata.source} to ${outFolder}`
diff --git a/collector/utils/extensions/RepoLoader/GitlabRepo/index.js b/collector/utils/extensions/RepoLoader/GitlabRepo/index.js
index f1c528f1d9f69d116e339286febe02be4f5d03b3..cd74fb3164bc0d1968522c8a0de1f09f95925569 100644
--- a/collector/utils/extensions/RepoLoader/GitlabRepo/index.js
+++ b/collector/utils/extensions/RepoLoader/GitlabRepo/index.js
@@ -82,7 +82,7 @@ async function loadGitlabRepo(args, response) {
}
data.wordCount = pageContent.split(" ").length;
- data.token_count_estimate = tokenizeString(pageContent).length;
+ data.token_count_estimate = tokenizeString(pageContent);
data.pageContent = pageContent;
console.log(
diff --git a/collector/utils/extensions/WebsiteDepth/index.js b/collector/utils/extensions/WebsiteDepth/index.js
index e680c0233b7b96cb46703757d9647ac5aa9ff613..4801a45aeba6d8095f33bcbc72ecf8dd6427621b 100644
--- a/collector/utils/extensions/WebsiteDepth/index.js
+++ b/collector/utils/extensions/WebsiteDepth/index.js
@@ -122,7 +122,7 @@ async function bulkScrapePages(links, outFolderPath) {
published: new Date().toLocaleString(),
wordCount: content.split(" ").length,
pageContent: content,
- token_count_estimate: tokenizeString(content).length,
+ token_count_estimate: tokenizeString(content),
};
writeToServerDocuments(data, data.title, outFolderPath);
diff --git a/collector/utils/extensions/YoutubeTranscript/index.js b/collector/utils/extensions/YoutubeTranscript/index.js
index c7cf7c1f836db7bbfe42acb97113f6930480b839..0e1e13feb140a8145016c1cee96117d03d4fc598 100644
--- a/collector/utils/extensions/YoutubeTranscript/index.js
+++ b/collector/utils/extensions/YoutubeTranscript/index.js
@@ -115,7 +115,7 @@ async function loadYouTubeTranscript({ url }) {
published: new Date().toLocaleString(),
wordCount: content.split(" ").length,
pageContent: content,
- token_count_estimate: tokenizeString(content).length,
+ token_count_estimate: tokenizeString(content),
};
console.log(`[YouTube Loader]: Saving ${metadata.title} to ${outFolder}`);
diff --git a/collector/utils/tokenizer/index.js b/collector/utils/tokenizer/index.js
index 618a7cdc7a6ceba013680461d0bf198061f162fe..2086be2574b55f18692b85c6fb31f75099ea8177 100644
--- a/collector/utils/tokenizer/index.js
+++ b/collector/utils/tokenizer/index.js
@@ -1,15 +1,66 @@
const { getEncoding } = require("js-tiktoken");
-function tokenizeString(input = "") {
- try {
- const encoder = getEncoding("cl100k_base");
- return encoder.encode(input);
- } catch (e) {
- console.error("Could not tokenize string!");
- return [];
+class TikTokenTokenizer {
+ static MAX_KB_ESTIMATE = 10;
+ static DIVISOR = 8;
+
+ constructor() {
+ if (TikTokenTokenizer.instance) {
+ this.log(
+ "Singleton instance already exists. Returning existing instance."
+ );
+ return TikTokenTokenizer.instance;
+ }
+
+ this.encoder = getEncoding("cl100k_base");
+ TikTokenTokenizer.instance = this;
+ this.log("Initialized new TikTokenTokenizer instance.");
+ }
+
+ log(text, ...args) {
+ console.log(`\x1b[35m[TikTokenTokenizer]\x1b[0m ${text}`, ...args);
+ }
+
+ /**
+ * Check if the input is too long to encode
+ * this is more of a rough estimate and a sanity check to prevent
+ * CPU issues from encoding too large of strings
+ * Assumes 1 character = 2 bytes in JS
+ * @param {string} input
+ * @returns {boolean}
+ */
+ #isTooLong(input) {
+ const bytesEstimate = input.length * 2;
+ const kbEstimate = Math.floor(bytesEstimate / 1024);
+ return kbEstimate >= TikTokenTokenizer.MAX_KB_ESTIMATE;
+ }
+
+ /**
+ * Encode a string into tokens for rough token count estimation.
+ * @param {string} input
+ * @returns {number}
+ */
+ tokenizeString(input = "") {
+ try {
+ if (this.#isTooLong(input)) {
+ this.log("Input will take too long to encode - estimating");
+ return Math.ceil(input.length / TikTokenTokenizer.DIVISOR);
+ }
+
+ return this.encoder.encode(input).length;
+ } catch (e) {
+ this.log("Could not tokenize string! Estimating...", e.message, e.stack);
+ return Math.ceil(input?.length / TikTokenTokenizer.DIVISOR) || 0;
+ }
}
}
+const tokenizer = new TikTokenTokenizer();
module.exports = {
- tokenizeString,
+ /**
+ * Encode a string into tokens for rough token count estimation.
+ * @param {string} input
+ * @returns {number}
+ */
+ tokenizeString: (input) => tokenizer.tokenizeString(input),
};
diff --git a/server/utils/AiProviders/deepseek/index.js b/server/utils/AiProviders/deepseek/index.js
index 694f85501c2a931cf7e2e034a4691c882191b897..b91332a84a15ae44bdb8e266b4527e78a4408462 100644
--- a/server/utils/AiProviders/deepseek/index.js
+++ b/server/utils/AiProviders/deepseek/index.js
@@ -4,7 +4,10 @@ const {
} = require("../../helpers/chat/LLMPerformanceMonitor");
const { v4: uuidv4 } = require("uuid");
const { MODEL_MAP } = require("../modelMap");
-const { writeResponseChunk, clientAbortedHandler } = require("../../helpers/chat/responses");
+const {
+ writeResponseChunk,
+ clientAbortedHandler,
+} = require("../../helpers/chat/responses");
class DeepSeekLLM {
constructor(embedder = null, modelPreference = null) {
diff --git a/server/utils/helpers/tiktoken.js b/server/utils/helpers/tiktoken.js
index a3fa3b63966411ef214ef1e979b0012d39137ec3..394f26187433d209664b0eef2b38b31acb6c59d6 100644
--- a/server/utils/helpers/tiktoken.js
+++ b/server/utils/helpers/tiktoken.js
@@ -1,10 +1,36 @@
const { getEncodingNameForModel, getEncoding } = require("js-tiktoken");
+/**
+ * @class TokenManager
+ *
+ * @notice
+ * We cannot do estimation of tokens here like we do in the collector
+ * because we need to know the model to do it.
+ * Other issues are we also do reverse tokenization here for the chat history during cannonballing.
+ * So here we are stuck doing the actual tokenization and encoding until we figure out what to do with prompt overflows.
+ */
class TokenManager {
+ static instance = null;
+ static currentModel = null;
+
constructor(model = "gpt-3.5-turbo") {
+ if (TokenManager.instance && TokenManager.currentModel === model) {
+ this.log("Returning existing instance for model:", model);
+ return TokenManager.instance;
+ }
+
this.model = model;
this.encoderName = this.#getEncodingFromModel(model);
this.encoder = getEncoding(this.encoderName);
+
+ TokenManager.instance = this;
+ TokenManager.currentModel = model;
+ this.log("Initialized new TokenManager instance for model:", model);
+ return this;
+ }
+
+ log(text, ...args) {
+ console.log(`\x1b[35m[TokenManager]\x1b[0m ${text}`, ...args);
}
#getEncodingFromModel(model) {
@@ -15,9 +41,11 @@ class TokenManager {
}
}
- // Pass in an empty array of disallowedSpecials to handle all tokens as text and to be tokenized.
- // https://github.com/openai/tiktoken/blob/9e79899bc248d5313c7dd73562b5e211d728723d/tiktoken/core.py#L91C20-L91C38
- // Returns number[]
+ /**
+ * Pass in an empty array of disallowedSpecials to handle all tokens as text and to be tokenized.
+ * @param {string} input
+ * @returns {number[]}
+ */
tokensFromString(input = "") {
try {
const tokens = this.encoder.encode(String(input), undefined, []);
@@ -28,17 +56,31 @@ class TokenManager {
}
}
+ /**
+ * Converts an array of tokens back to a string.
+ * @param {number[]} tokens
+ * @returns {string}
+ */
bytesFromTokens(tokens = []) {
const bytes = this.encoder.decode(tokens);
return bytes;
}
- // Returns number
+ /**
+ * Counts the number of tokens in a string.
+ * @param {string} input
+ * @returns {number}
+ */
countFromString(input = "") {
const tokens = this.tokensFromString(input);
return tokens.length;
}
+ /**
+ * Estimates the number of tokens in a string or array of strings.
+ * @param {string | string[]} input
+ * @returns {number}
+ */
statsFrom(input) {
if (typeof input === "string") return this.countFromString(input);