@mastra/rag 1.0.6 → 1.0.7-alpha.0

package/src/document/document.ts

@@ -7,8 +7,23 @@ import { HTMLHeaderTransformer, HTMLSectionTransformer } from './transformers/ht
 import { RecursiveJsonTransformer } from './transformers/json';
 import { LatexTransformer } from './transformers/latex';
 import { MarkdownHeaderTransformer, MarkdownTransformer } from './transformers/markdown';
+import { SentenceTransformer } from './transformers/sentence';
 import { TokenTransformer } from './transformers/token';
-import type { ChunkOptions, ChunkParams, ChunkStrategy, ExtractParams } from './types';
+import type {
+  ChunkParams,
+  ChunkStrategy,
+  ExtractParams,
+  HTMLChunkOptions,
+  RecursiveChunkOptions,
+  CharacterChunkOptions,
+  TokenChunkOptions,
+  MarkdownChunkOptions,
+  JsonChunkOptions,
+  LatexChunkOptions,
+  SentenceChunkOptions,
+  StrategyOptions,
+} from './types';
+import { validateChunkParams } from './validation';
 
 export class MDocument {
   private chunks: Chunk[];
@@ -135,35 +150,27 @@ export class MDocument {
     }
   }
 
-  private async chunkBy(strategy: ChunkStrategy, options?: ChunkOptions): Promise<void> {
-    switch (strategy) {
-      case 'recursive':
-        await this.chunkRecursive(options);
-        break;
-      case 'character':
-        await this.chunkCharacter(options);
-        break;
-      case 'token':
-        await this.chunkToken(options);
-        break;
-      case 'markdown':
-        await this.chunkMarkdown(options);
-        break;
-      case 'html':
-        await this.chunkHTML(options);
-        break;
-      case 'json':
-        await this.chunkJSON(options);
-        break;
-      case 'latex':
-        await this.chunkLatex(options);
-        break;
-      default:
-        throw new Error(`Unknown strategy: ${strategy}`);
+  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];
+    if (chunkingFunc) {
+      await chunkingFunc(options);
+    } else {
+      throw new Error(`Unknown strategy: ${strategy}`);
     }
   }
 
-  async chunkRecursive(options?: ChunkOptions): Promise<void> {
+  async chunkRecursive(options?: RecursiveChunkOptions): Promise<void> {
     if (options?.language) {
       const rt = RecursiveCharacterTransformer.fromLanguage(options.language, options);
       const textSplit = rt.transformDocuments(this.chunks);
@@ -171,28 +178,24 @@ export class MDocument {
       return;
     }
 
-    const rt = new RecursiveCharacterTransformer({
-      separators: options?.separators,
-      isSeparatorRegex: options?.isSeparatorRegex,
-      options,
-    });
+    const rt = new RecursiveCharacterTransformer(options);
     const textSplit = rt.transformDocuments(this.chunks);
     this.chunks = textSplit;
   }
 
-  async chunkCharacter(options?: ChunkOptions): Promise<void> {
+  async chunkCharacter(options?: CharacterChunkOptions): Promise<void> {
     const rt = new CharacterTransformer({
+      ...options,
       separator: options?.separator,
       isSeparatorRegex: options?.isSeparatorRegex,
-      options,
     });
     const textSplit = rt.transformDocuments(this.chunks);
     this.chunks = textSplit;
   }
 
-  async chunkHTML(options?: ChunkOptions): Promise<void> {
+  async chunkHTML(options?: HTMLChunkOptions): Promise<void> {
     if (options?.headers?.length) {
-      const rt = new HTMLHeaderTransformer(options.headers, options?.returnEachLine);
+      const rt = new HTMLHeaderTransformer(options as HTMLChunkOptions & { headers: [string, string][] });
 
       const textSplit = rt.transformDocuments(this.chunks);
       this.chunks = textSplit;
@@ -200,7 +203,7 @@ export class MDocument {
     }
 
     if (options?.sections?.length) {
-      const rt = new HTMLSectionTransformer(options.sections);
+      const rt = new HTMLSectionTransformer(options as HTMLChunkOptions & { sections: [string, string][] });
 
       const textSplit = rt.transformDocuments(this.chunks);
       this.chunks = textSplit;
@@ -210,7 +213,7 @@ export class MDocument {
     throw new Error('HTML chunking requires either headers or sections to be specified');
   }
 
-  async chunkJSON(options?: ChunkOptions): Promise<void> {
+  async chunkJSON(options?: JsonChunkOptions): Promise<void> {
     if (!options?.maxSize) {
       throw new Error('JSON chunking requires maxSize to be specified');
     }
@@ -229,13 +232,13 @@ export class MDocument {
     this.chunks = textSplit;
   }
 
-  async chunkLatex(options?: ChunkOptions): Promise<void> {
+  async chunkLatex(options?: LatexChunkOptions): Promise<void> {
     const rt = new LatexTransformer(options);
     const textSplit = rt.transformDocuments(this.chunks);
     this.chunks = textSplit;
   }
 
-  async chunkToken(options?: ChunkOptions): Promise<void> {
+  async chunkToken(options?: TokenChunkOptions): Promise<void> {
     const rt = TokenTransformer.fromTikToken({
       options,
       encodingName: options?.encodingName,
@@ -245,7 +248,7 @@ export class MDocument {
     this.chunks = textSplit;
   }
 
-  async chunkMarkdown(options?: ChunkOptions): Promise<void> {
+  async chunkMarkdown(options?: MarkdownChunkOptions): Promise<void> {
     if (options?.headers) {
       const rt = new MarkdownHeaderTransformer(options.headers, options?.returnEachLine, options?.stripHeaders);
       const textSplit = rt.transformDocuments(this.chunks);
@@ -258,11 +261,36 @@ export class MDocument {
     this.chunks = textSplit;
   }
 
+  async chunkSentence(options?: SentenceChunkOptions): Promise<void> {
+    if (!options?.maxSize) {
+      throw new Error('Sentence chunking requires maxSize to be specified');
+    }
+
+    const rt = new SentenceTransformer({
+      minSize: options?.minSize,
+      maxSize: options?.maxSize,
+      targetSize: options?.targetSize,
+      overlap: options?.overlap,
+      sentenceEnders: options?.sentenceEnders,
+      fallbackToWords: options?.fallbackToWords,
+      fallbackToCharacters: options?.fallbackToCharacters,
+      keepSeparator: options?.keepSeparator,
+      lengthFunction: options?.lengthFunction,
+      addStartIndex: options?.addStartIndex,
+      stripWhitespace: options?.stripWhitespace,
+    });
+
+    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
     const strategy = passedStrategy || this.defaultStrategy();
 
+    validateChunkParams(strategy, chunkOptions);
+
     // Apply the appropriate chunking strategy
     await this.chunkBy(strategy, chunkOptions);
 

package/src/document/transformers/character.ts

@@ -1,5 +1,5 @@
 import { Language } from '../types';
-import type { ChunkOptions } from '../types';
+import type { BaseChunkOptions, CharacterChunkOptions, RecursiveChunkOptions } from '../types';
 
 import { TextTransformer } from './text';
 
@@ -52,23 +52,8 @@ export class CharacterTransformer extends TextTransformer {
   protected separator: string;
   protected isSeparatorRegex: boolean;
 
-  constructor({
-    separator = '\n\n',
-    isSeparatorRegex = false,
-    options = {},
-  }: {
-    separator?: string;
-    isSeparatorRegex?: boolean;
-    options?: {
-      size?: number;
-      overlap?: number;
-      lengthFunction?: (text: string) => number;
-      keepSeparator?: boolean | 'start' | 'end';
-      addStartIndex?: boolean;
-      stripWhitespace?: boolean;
-    };
-  }) {
-    super(options);
+  constructor({ separator = '\n\n', isSeparatorRegex = false, ...baseOptions }: CharacterChunkOptions = {}) {
+    super(baseOptions);
     this.separator = separator;
     this.isSeparatorRegex = isSeparatorRegex;
   }
@@ -82,7 +67,7 @@ export class CharacterTransformer extends TextTransformer {
     // If length of any split is greater than chunk size, perform additional splitting
     const chunks: string[] = [];
     for (const split of initialSplits) {
-      if (this.lengthFunction(split) <= this.size) {
+      if (this.lengthFunction(split) <= this.maxSize) {
         chunks.push(split);
       } else {
         // If a single split is too large, split it further with overlap
@@ -102,7 +87,7 @@ export class CharacterTransformer extends TextTransformer {
       let chunkEnd = currentPosition;
 
       // Build chunk up to max size
-      while (chunkEnd < text.length && this.lengthFunction(text.slice(currentPosition, chunkEnd + 1)) <= this.size) {
+      while (chunkEnd < text.length && this.lengthFunction(text.slice(currentPosition, chunkEnd + 1)) <= this.maxSize) {
         chunkEnd++;
       }
 
@@ -125,16 +110,8 @@ export class RecursiveCharacterTransformer extends TextTransformer {
   protected separators: string[];
   protected isSeparatorRegex: boolean;
 
-  constructor({
-    separators,
-    isSeparatorRegex = false,
-    options = {},
-  }: {
-    separators?: string[];
-    isSeparatorRegex?: boolean;
-    options?: ChunkOptions;
-  }) {
-    super(options);
+  constructor({ separators, isSeparatorRegex = false, language, ...baseOptions }: RecursiveChunkOptions = {}) {
+    super(baseOptions);
     this.separators = separators || ['\n\n', '\n', ' ', ''];
     this.isSeparatorRegex = isSeparatorRegex;
   }
@@ -169,7 +146,7 @@ export class RecursiveCharacterTransformer extends TextTransformer {
     const mergeSeparator = this.keepSeparator ? '' : separator;
 
     for (const s of splits) {
-      if (this.lengthFunction(s) < this.size) {
+      if (this.lengthFunction(s) < this.maxSize) {
         goodSplits.push(s);
       } else {
         if (goodSplits.length > 0) {
@@ -198,19 +175,14 @@ export class RecursiveCharacterTransformer extends TextTransformer {
     return this._splitText(text, this.separators);
   }
 
-  static fromLanguage(
-    language: Language,
-    options: {
-      size?: number;
-      chunkOverlap?: number;
-      lengthFunction?: (text: string) => number;
-      keepSeparator?: boolean | 'start' | 'end';
-      addStartIndex?: boolean;
-      stripWhitespace?: boolean;
-    } = {},
-  ): RecursiveCharacterTransformer {
+  static fromLanguage(language: Language, options: BaseChunkOptions = {}): RecursiveCharacterTransformer {
     const separators = RecursiveCharacterTransformer.getSeparatorsForLanguage(language);
-    return new RecursiveCharacterTransformer({ separators, isSeparatorRegex: true, options });
+    return new RecursiveCharacterTransformer({
+      ...options,
+      separators,
+      isSeparatorRegex: true,
+      language,
+    });
   }
 
   static getSeparatorsForLanguage(language: Language): string[] {

package/src/document/transformers/html.ts

@@ -1,5 +1,6 @@
 import { parse } from 'node-html-better-parser';
 import { Document } from '../schema';
+import type { HTMLChunkOptions } from '../types';
 
 import { RecursiveCharacterTransformer } from './character';
 
@@ -14,9 +15,9 @@ export class HTMLHeaderTransformer {
   private headersToSplitOn: [string, string][];
   private returnEachElement: boolean;
 
-  constructor(headersToSplitOn: [string, string][], returnEachElement: boolean = false) {
-    this.returnEachElement = returnEachElement;
-    this.headersToSplitOn = [...headersToSplitOn].sort();
+  constructor(options: HTMLChunkOptions & { headers: [string, string][] }) {
+    this.returnEachElement = options.returnEachLine ?? false;
+    this.headersToSplitOn = [...options.headers].sort();
   }
 
   splitText({ text }: { text: string }): Document[] {
@@ -195,11 +196,11 @@ export class HTMLHeaderTransformer {
 
 export class HTMLSectionTransformer {
   private headersToSplitOn: Record<string, string>;
-  private options: Record<string, any>;
+  private textSplitter: RecursiveCharacterTransformer;
 
-  constructor(headersToSplitOn: [string, string][], options: Record<string, any> = {}) {
-    this.headersToSplitOn = Object.fromEntries(headersToSplitOn.map(([tag, name]) => [tag.toLowerCase(), name]));
-    this.options = options;
+  constructor(options: HTMLChunkOptions & { sections: [string, string][] }) {
+    this.headersToSplitOn = Object.fromEntries(options.sections.map(([tag, name]) => [tag.toLowerCase(), name]));
+    this.textSplitter = new RecursiveCharacterTransformer(options);
   }
 
   splitText(text: string): Document[] {
@@ -296,9 +297,8 @@ export class HTMLSectionTransformer {
       metadatas.push(doc.metadata);
     }
     const results = await this.createDocuments(texts, metadatas);
-    const textSplitter = new RecursiveCharacterTransformer({ options: this.options });
 
-    return textSplitter.splitDocuments(results);
+    return this.textSplitter.splitDocuments(results);
   }
 
   createDocuments(texts: string[], metadatas?: Record<string, any>[]): Document[] {

package/src/document/transformers/json.ts

@@ -1,12 +1,17 @@
 import { Document } from '../schema';
+import type { JsonChunkOptions } from '../types';
 
 export class RecursiveJsonTransformer {
   private maxSize: number;
   private minSize: number;
+  private ensureAscii: boolean;
+  private convertLists: boolean;
 
-  constructor({ maxSize = 2000, minSize }: { maxSize: number; minSize?: number }) {
+  constructor({ maxSize = 2000, minSize, ensureAscii = false, convertLists = true }: JsonChunkOptions) {
     this.maxSize = maxSize;
     this.minSize = minSize ?? Math.max(maxSize - 200, 50);
+    this.ensureAscii = ensureAscii;
+    this.convertLists = convertLists;
   }
 
   private static jsonSize(data: Record<string, any>): number {
@@ -170,8 +175,8 @@ export class RecursiveJsonTransformer {
   private isWithinSizeLimit(value: any, currentSize: number = 0): boolean {
     const size = RecursiveJsonTransformer.jsonSize(value);
     // If this is a new chunk (currentSize = 0), allow items smaller than maxSize
-    // If adding to existing chunk, ensure we're above minSize before splitting
-    return currentSize === 0 ? size <= this.maxSize : size + currentSize <= this.maxSize || currentSize < this.minSize;
+    // If adding to existing chunk, ensure total size doesn't exceed maxSize
+    return currentSize === 0 ? size <= this.maxSize : size + currentSize <= this.maxSize;
   }
 
   /**

package/src/document/transformers/latex.ts

@@ -1,19 +1,11 @@
 import { Language } from '../types';
+import type { BaseChunkOptions } from '../types';
 
 import { RecursiveCharacterTransformer } from './character';
 
 export class LatexTransformer extends RecursiveCharacterTransformer {
-  constructor(
-    options: {
-      size?: number;
-      overlap?: number;
-      lengthFunction?: (text: string) => number;
-      keepSeparator?: boolean | 'start' | 'end';
-      addStartIndex?: boolean;
-      stripWhitespace?: boolean;
-    } = {},
-  ) {
+  constructor(options: BaseChunkOptions = {}) {
     const separators = RecursiveCharacterTransformer.getSeparatorsForLanguage(Language.LATEX);
-    super({ separators, isSeparatorRegex: true, options });
+    super({ ...options, separators, isSeparatorRegex: true });
   }
 }

package/src/document/transformers/markdown.ts

@@ -1,6 +1,7 @@
 import { Document } from '../schema';
 
 import { Language } from '../types';
+import type { BaseChunkOptions } from '../types';
 
 import { RecursiveCharacterTransformer } from './character';
 
@@ -16,18 +17,9 @@ interface HeaderType {
 }
 
 export class MarkdownTransformer extends RecursiveCharacterTransformer {
-  constructor(
-    options: {
-      chunkSize?: number;
-      chunkOverlap?: number;
-      lengthFunction?: (text: string) => number;
-      keepSeparator?: boolean | 'start' | 'end';
-      addStartIndex?: boolean;
-      stripWhitespace?: boolean;
-    } = {},
-  ) {
+  constructor(options: BaseChunkOptions = {}) {
     const separators = RecursiveCharacterTransformer.getSeparatorsForLanguage(Language.MARKDOWN);
-    super({ separators, isSeparatorRegex: true, options });
+    super({ ...options, separators, isSeparatorRegex: true });
   }
 }