@mastra/rag 2.1.1 → 2.1.2

package/dist/index.js

@@ -6066,31 +6066,34 @@ var MarkdownHeaderTransformer = class {
 var SemanticMarkdownTransformer = class _SemanticMarkdownTransformer extends TextTransformer {
   tokenizer;
   joinThreshold;
-  allowedSpecial;
-  disallowedSpecial;
+  allowedArray;
+  disallowedArray;
   constructor({
     joinThreshold = 500,
     encodingName = "cl100k_base",
     modelName,
+    tokenizer: existingTokenizer,
     allowedSpecial = /* @__PURE__ */ new Set(),
     disallowedSpecial = "all",
     ...baseOptions
   } = {}) {
     super(baseOptions);
     this.joinThreshold = joinThreshold;
-    this.allowedSpecial = allowedSpecial;
-    this.disallowedSpecial = disallowedSpecial;
-    try {
-      this.tokenizer = modelName ? encodingForModel(modelName) : getEncoding(encodingName);
-    } catch {
-      throw new Error("Could not load tiktoken encoding. Please install it with `npm install js-tiktoken`.");
+    this.allowedArray = allowedSpecial === "all" ? "all" : Array.from(allowedSpecial);
+    this.disallowedArray = disallowedSpecial === "all" ? "all" : Array.from(disallowedSpecial);
+    if (existingTokenizer) {
+      this.tokenizer = existingTokenizer;
+    } else {
+      try {
+        this.tokenizer = modelName ? encodingForModel(modelName) : getEncoding(encodingName);
+      } catch {
+        throw new Error("Could not load tiktoken encoding. Please install it with `npm install js-tiktoken`.");
+      }
     }
   }
   countTokens(text) {
-    const allowed = this.allowedSpecial === "all" ? "all" : Array.from(this.allowedSpecial);
-    const disallowed = this.disallowedSpecial === "all" ? "all" : Array.from(this.disallowedSpecial);
     const processedText = this.stripWhitespace ? text.trim() : text;
-    return this.tokenizer.encode(processedText, allowed, disallowed).length;
+    return this.tokenizer.encode(processedText, this.allowedArray, this.disallowedArray).length;
   }
   splitMarkdownByHeaders(markdown) {
     const sections = [];
@@ -6144,14 +6147,21 @@ var SemanticMarkdownTransformer = class _SemanticMarkdownTransformer extends Tex
         const current = workingSections[j];
         if (current.depth === depth) {
           const prev = workingSections[j - 1];
-          if (prev.length + current.length < this.joinThreshold && prev.depth <= current.depth) {
-            const title = `${"#".repeat(current.depth)} ${current.title}`;
-            const formattedTitle = `
+          const title = `${"#".repeat(current.depth)} ${current.title}`;
+          const formattedTitle = `
 
 ${title}`;
+          const headerLength = this.tokenizer.encode(
+            `${formattedTitle}
+`,
+            this.allowedArray,
+            this.disallowedArray
+          ).length;
+          const mergedLength = prev.length + current.length + headerLength;
+          if (mergedLength < this.joinThreshold && prev.depth <= current.depth) {
             prev.content += `${formattedTitle}
 ${current.content}`;
-            prev.length = this.countTokens(prev.content);
+            prev.length = mergedLength;
             workingSections.splice(j, 1);
             j--;
           }
@@ -6221,6 +6231,7 @@ ${section.content}`;
       ...options,
       encodingName,
       modelName,
+      tokenizer,
       lengthFunction: tikTokenCounter
     });
   }
@@ -6475,30 +6486,33 @@ function splitTextOnTokens({ text, tokenizer }) {
 }
 var TokenTransformer = class _TokenTransformer extends TextTransformer {
   tokenizer;
-  allowedSpecial;
-  disallowedSpecial;
+  allowedArray;
+  disallowedArray;
   constructor({
     encodingName = "cl100k_base",
     modelName,
+    tokenizer: existingTokenizer,
     allowedSpecial = /* @__PURE__ */ new Set(),
     disallowedSpecial = "all",
     options = {}
   }) {
     super(options);
-    try {
-      this.tokenizer = modelName ? encodingForModel(modelName) : getEncoding(encodingName);
-    } catch {
-      throw new Error("Could not load tiktoken encoding. Please install it with `npm install js-tiktoken`.");
+    if (existingTokenizer) {
+      this.tokenizer = existingTokenizer;
+    } else {
+      try {
+        this.tokenizer = modelName ? encodingForModel(modelName) : getEncoding(encodingName);
+      } catch {
+        throw new Error("Could not load tiktoken encoding. Please install it with `npm install js-tiktoken`.");
+      }
     }
-    this.allowedSpecial = allowedSpecial;
-    this.disallowedSpecial = disallowedSpecial;
+    this.allowedArray = allowedSpecial === "all" ? "all" : Array.from(allowedSpecial);
+    this.disallowedArray = disallowedSpecial === "all" ? "all" : Array.from(disallowedSpecial);
   }
   splitText({ text }) {
     const encode = (text2) => {
-      const allowed = this.allowedSpecial === "all" ? "all" : Array.from(this.allowedSpecial);
-      const disallowed = this.disallowedSpecial === "all" ? "all" : Array.from(this.disallowedSpecial);
       const processedText = this.stripWhitespace ? text2.trim() : text2;
-      return Array.from(this.tokenizer.encode(processedText, allowed, disallowed));
+      return Array.from(this.tokenizer.encode(processedText, this.allowedArray, this.disallowedArray));
     };
     const decode = (tokens) => {
       const text2 = this.tokenizer.decode(tokens);
@@ -6535,6 +6549,7 @@ var TokenTransformer = class _TokenTransformer extends TextTransformer {
     return new _TokenTransformer({
       encodingName,
       modelName,
+      tokenizer,
       allowedSpecial: options.allowedSpecial,
       disallowedSpecial: options.disallowedSpecial,
       options: {