@mastra/rag 1.0.7 → 1.0.8-alpha.0

package/src/document/document.ts

@@ -7,6 +7,7 @@ import { HTMLHeaderTransformer, HTMLSectionTransformer } from './transformers/ht
 import { RecursiveJsonTransformer } from './transformers/json';
 import { LatexTransformer } from './transformers/latex';
 import { MarkdownHeaderTransformer, MarkdownTransformer } from './transformers/markdown';
+import { SemanticMarkdownTransformer } from './transformers/semantic-markdown';
 import { SentenceTransformer } from './transformers/sentence';
 import { TokenTransformer } from './transformers/token';
 import type {
@@ -18,6 +19,7 @@ import type {
   CharacterChunkOptions,
   TokenChunkOptions,
   MarkdownChunkOptions,
+  SemanticMarkdownChunkOptions,
   JsonChunkOptions,
   LatexChunkOptions,
   SentenceChunkOptions,
@@ -150,19 +152,27 @@ export class MDocument {
     }
   }
 
+  private _strategyMap?: { [S in ChunkStrategy]: (options?: StrategyOptions[S]) => Promise<void> };
+
+  private get strategyMap() {
+    if (!this._strategyMap) {
+      this._strategyMap = {
+        recursive: options => this.chunkRecursive(options),
+        character: options => this.chunkCharacter(options),
+        token: options => this.chunkToken(options),
+        markdown: options => this.chunkMarkdown(options),
+        html: options => this.chunkHTML(options),
+        json: options => this.chunkJSON(options),
+        latex: options => this.chunkLatex(options),
+        sentence: options => this.chunkSentence(options),
+        'semantic-markdown': options => this.chunkSemanticMarkdown(options),
+      };
+    }
+    return this._strategyMap;
+  }
+
   private async chunkBy<K extends ChunkStrategy>(strategy: K, options?: StrategyOptions[K]): Promise<void> {
-    const strategyMap: { [S in ChunkStrategy]: (options?: StrategyOptions[S]) => Promise<void> } = {
-      recursive: options => this.chunkRecursive(options),
-      character: options => this.chunkCharacter(options),
-      token: options => this.chunkToken(options),
-      markdown: options => this.chunkMarkdown(options),
-      html: options => this.chunkHTML(options),
-      json: options => this.chunkJSON(options),
-      latex: options => this.chunkLatex(options),
-      sentence: options => this.chunkSentence(options),
-    };
-
-    const chunkingFunc = strategyMap[strategy];
+    const chunkingFunc = this.strategyMap[strategy];
     if (chunkingFunc) {
       await chunkingFunc(options);
     } else {
@@ -284,6 +294,16 @@ export class MDocument {
     this.chunks = textSplit;
   }
 
+  async chunkSemanticMarkdown(options?: SemanticMarkdownChunkOptions): Promise<void> {
+    const rt = SemanticMarkdownTransformer.fromTikToken({
+      options,
+      encodingName: options?.encodingName,
+      modelName: options?.modelName,
+    });
+    const textSplit = rt.transformDocuments(this.chunks);
+    this.chunks = textSplit;
+  }
+
   async chunk(params?: ChunkParams): Promise<Chunk[]> {
     const { strategy: passedStrategy, extract, ...chunkOptions } = params || {};
     // Determine the default strategy based on type if not specified

package/src/document/transformers/semantic-markdown.ts

@@ -0,0 +1,227 @@
+import type { TiktokenModel, TiktokenEncoding, Tiktoken } from 'js-tiktoken';
+import { encodingForModel, getEncoding } from 'js-tiktoken';
+import { Document } from '../schema';
+import type { SemanticMarkdownChunkOptions } from '../types';
+
+import { TextTransformer } from './text';
+
+interface MarkdownNode {
+  title: string;
+  depth: number;
+  content: string;
+  length: number;
+}
+
+export class SemanticMarkdownTransformer extends TextTransformer {
+  private tokenizer: Tiktoken;
+  private joinThreshold: number;
+  private allowedSpecial: Set<string> | 'all';
+  private disallowedSpecial: Set<string> | 'all';
+
+  constructor({
+    joinThreshold = 500,
+    encodingName = 'cl100k_base',
+    modelName,
+    allowedSpecial = new Set(),
+    disallowedSpecial = 'all',
+    ...baseOptions
+  }: SemanticMarkdownChunkOptions = {}) {
+    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`.');
+    }
+  }
+
+  private countTokens(text: string): number {
+    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;
+  }
+
+  private splitMarkdownByHeaders(markdown: string): MarkdownNode[] {
+    const sections: MarkdownNode[] = [];
+    const lines = markdown.split('\n');
+    let currentContent = '';
+    let currentTitle = '';
+    let currentDepth = 0;
+    let inCodeBlock = false;
+
+    const headerRegex = /^(#+)\s+(.+)$/;
+
+    for (let i = 0; i < lines.length; i++) {
+      const line = lines[i]!;
+      const headerMatch = line.match(headerRegex);
+
+      // Track code blocks to avoid parsing headers inside them
+      if (line.startsWith('```') || line.startsWith('~~~')) {
+        inCodeBlock = !inCodeBlock;
+      }
+
+      if (headerMatch && !inCodeBlock) {
+        // Save previous section
+        // Push the previous section if it has content or if it's a header.
+        // This ensures headers that only act as parents are not lost.
+        if (currentContent.trim() !== '' || (currentTitle && currentDepth > 0)) {
+          sections.push({
+            title: currentTitle,
+            content: currentContent.trim(),
+            depth: currentDepth,
+            length: this.countTokens(currentContent.trim()),
+          });
+        }
+        currentContent = ''; // Always reset for the new section
+
+        // Start new section
+        currentDepth = headerMatch[1]!.length;
+        currentTitle = headerMatch[2]!;
+      } else {
+        currentContent += line + '\n';
+      }
+    }
+
+    // Add the last section
+    if (currentContent.trim() !== '') {
+      sections.push({
+        title: currentTitle,
+        content: currentContent.trim(),
+        depth: currentDepth,
+        length: this.countTokens(currentContent.trim()),
+      });
+    }
+
+    // Remove initial empty preamble if present, but keep non-empty preambles
+    if (sections.length > 1 && sections[0]!.title === '' && sections[0]!.content.trim() === '') {
+      sections.shift();
+    }
+
+    return sections;
+  }
+
+  private mergeSemanticSections(sections: MarkdownNode[]): MarkdownNode[] {
+    if (sections.length === 0) return sections;
+
+    const workingSections = [...sections];
+    const deepest = Math.max(...workingSections.map(s => s.depth));
+
+    for (let depth = deepest; depth > 0; depth--) {
+      for (let j = 1; j < workingSections.length; j++) {
+        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 = `\n\n${title}`;
+
+            prev.content += `${formattedTitle}\n${current.content}`;
+
+            prev.length = this.countTokens(prev.content);
+
+            workingSections.splice(j, 1);
+            j--;
+          }
+        }
+      }
+    }
+
+    return workingSections;
+  }
+
+  splitText({ text }: { text: string }): string[] {
+    if (!text.trim()) return [];
+
+    const initialSections = this.splitMarkdownByHeaders(text);
+
+    const mergedSections = this.mergeSemanticSections(initialSections);
+
+    return mergedSections.map(section => {
+      if (section.title) {
+        const header = `${'#'.repeat(section.depth)} ${section.title}`;
+        return `${header}\n${section.content}`;
+      }
+      return section.content;
+    });
+  }
+
+  createDocuments(texts: string[], metadatas?: Record<string, any>[]): Document[] {
+    const _metadatas = metadatas || Array(texts.length).fill({});
+    const documents: Document[] = [];
+
+    texts.forEach((text, i) => {
+      this.splitText({ text }).forEach(chunk => {
+        const metadata = {
+          ..._metadatas[i],
+          tokenCount: this.countTokens(chunk),
+        };
+
+        documents.push(
+          new Document({
+            text: chunk,
+            metadata,
+          }),
+        );
+      });
+    });
+
+    return documents;
+  }
+
+  transformDocuments(documents: Document[]): Document[] {
+    const texts: string[] = [];
+    const metadatas: Record<string, any>[] = [];
+
+    for (const doc of documents) {
+      texts.push(doc.text);
+      metadatas.push(doc.metadata);
+    }
+
+    return this.createDocuments(texts, metadatas);
+  }
+
+  static fromTikToken({
+    encodingName = 'cl100k_base',
+    modelName,
+    options = {},
+  }: {
+    encodingName?: TiktokenEncoding;
+    modelName?: TiktokenModel;
+    options?: SemanticMarkdownChunkOptions;
+  }): SemanticMarkdownTransformer {
+    let tokenizer: Tiktoken;
+
+    try {
+      tokenizer = modelName ? encodingForModel(modelName) : getEncoding(encodingName);
+    } catch {
+      throw new Error('Could not load tiktoken encoding. Please install it with `npm install js-tiktoken`.');
+    }
+
+    const tikTokenCounter = (text: string): number => {
+      const allowed =
+        options.allowedSpecial === 'all' ? 'all' : options.allowedSpecial ? Array.from(options.allowedSpecial) : [];
+      const disallowed =
+        options.disallowedSpecial === 'all'
+          ? 'all'
+          : options.disallowedSpecial
+            ? Array.from(options.disallowedSpecial)
+            : [];
+      return tokenizer.encode(text, allowed, disallowed).length;
+    };
+
+    return new SemanticMarkdownTransformer({
+      ...options,
+      encodingName,
+      modelName,
+      lengthFunction: tikTokenCounter,
+    });
+  }
+}

package/src/document/types.ts

@@ -79,6 +79,14 @@ export type MarkdownChunkOptions = BaseChunkOptions & {
   stripHeaders?: boolean;
 };
 
+export type SemanticMarkdownChunkOptions = BaseChunkOptions & {
+  joinThreshold?: number;
+  encodingName?: TiktokenEncoding;
+  modelName?: TiktokenModel;
+  allowedSpecial?: Set<string> | 'all';
+  disallowedSpecial?: Set<string> | 'all';
+};
+
 export type HTMLChunkOptions = BaseChunkOptions &
   (
     | { headers: [string, string][]; sections?: never; returnEachLine?: boolean }
@@ -111,9 +119,19 @@ export type StrategyOptions = {
   json: JsonChunkOptions;
   latex: LatexChunkOptions;
   sentence: SentenceChunkOptions;
+  'semantic-markdown': SemanticMarkdownChunkOptions;
 };
 
-export type ChunkStrategy = 'recursive' | 'character' | 'token' | 'markdown' | 'html' | 'json' | 'latex' | 'sentence';
+export type ChunkStrategy =
+  | 'recursive'
+  | 'character'
+  | 'token'
+  | 'markdown'
+  | 'html'
+  | 'json'
+  | 'latex'
+  | 'sentence'
+  | 'semantic-markdown';
 
 export type ChunkParams =
   | ({ strategy?: 'character' } & CharacterChunkOptions & { extract?: ExtractParams })
@@ -123,4 +141,5 @@ export type ChunkParams =
   | ({ strategy: 'html' } & HTMLChunkOptions & { extract?: ExtractParams })
   | ({ strategy: 'json' } & JsonChunkOptions & { extract?: ExtractParams })
   | ({ strategy: 'latex' } & LatexChunkOptions & { extract?: ExtractParams })
-  | ({ strategy: 'sentence' } & SentenceChunkOptions & { extract?: ExtractParams });
+  | ({ strategy: 'sentence' } & SentenceChunkOptions & { extract?: ExtractParams })
+  | ({ strategy: 'semantic-markdown' } & SemanticMarkdownChunkOptions & { extract?: ExtractParams });

package/src/document/validation.ts

@@ -108,6 +108,16 @@ const markdownChunkOptionsSchema = baseChunkOptionsSchema
   })
   .strict();
 
+const semanticMarkdownChunkOptionsSchema = baseChunkOptionsSchema
+  .extend({
+    joinThreshold: z.number().positive().optional(),
+    encodingName: z.string().optional(),
+    modelName: z.string().optional(),
+    allowedSpecial: setOrAllSchema,
+    disallowedSpecial: setOrAllSchema,
+  })
+  .strict();
+
 const latexChunkOptionsSchema = baseChunkOptionsSchema.strict();
 
 // Strategy-specific validation schemas
@@ -119,6 +129,7 @@ const validationSchemas = {
   json: jsonChunkOptionsSchema.transform(handleDeprecatedSize),
   html: htmlChunkOptionsSchema.transform(handleDeprecatedSize),
   markdown: markdownChunkOptionsSchema.transform(handleDeprecatedSize),
+  'semantic-markdown': semanticMarkdownChunkOptionsSchema.transform(handleDeprecatedSize),
   latex: latexChunkOptionsSchema.transform(handleDeprecatedSize),
 } as const;