@mastra/rag 1.0.6 → 1.0.7-alpha.0

package/src/document/transformers/sentence.ts

@@ -0,0 +1,314 @@
+import type { SentenceChunkOptions } from '../types';
+import { TextTransformer } from './text';
+
+export class SentenceTransformer extends TextTransformer {
+  protected minSize: number;
+  protected maxSize: number;
+  protected targetSize: number;
+  protected sentenceEnders: string[];
+  protected fallbackToWords: boolean;
+  protected fallbackToCharacters: boolean;
+  protected keepSeparator: boolean | 'start' | 'end';
+
+  constructor(options: SentenceChunkOptions) {
+    // Ensure overlap doesn't exceed maxSize for parent validation
+    const parentOverlap = Math.min(options.overlap ?? 0, options.maxSize - 1);
+
+    const baseOptions = {
+      ...options,
+      overlap: parentOverlap, // Use adjusted overlap for parent
+    };
+
+    super(baseOptions);
+
+    this.maxSize = options.maxSize;
+    this.minSize = options.minSize ?? 50;
+    this.targetSize = options.targetSize ?? Math.floor(options.maxSize * 0.8);
+    this.sentenceEnders = options.sentenceEnders ?? ['.', '!', '?'];
+    this.fallbackToWords = options.fallbackToWords ?? true;
+    this.fallbackToCharacters = options.fallbackToCharacters ?? true;
+    this.keepSeparator = options.keepSeparator ?? false;
+
+    // Override with original overlap for our sentence logic
+    this.overlap = options.overlap ?? 0;
+  }
+
+  private detectSentenceBoundaries(text: string): string[] {
+    if (!text) return [];
+
+    const sentences: string[] = [];
+    let currentSentence = '';
+    let i = 0;
+
+    while (i < text.length) {
+      const char = text[i];
+      if (!char) break; // Safety check
+
+      currentSentence += char;
+
+      if (this.sentenceEnders.includes(char)) {
+        const remainingText = text.slice(i + 1);
+
+        if (this.isRealSentenceBoundary(currentSentence, remainingText)) {
+          sentences.push(currentSentence.trim());
+          currentSentence = '';
+        }
+      }
+      i++;
+    }
+
+    if (currentSentence.trim()) {
+      sentences.push(currentSentence.trim());
+    }
+
+    return sentences.filter(s => s.length > 0);
+  }
+
+  private isRealSentenceBoundary(currentSentence: string, remainingText: string): boolean {
+    if (!remainingText.trim()) {
+      return true;
+    }
+
+    if (!/^\s+[A-Z]/.test(remainingText)) {
+      return false;
+    }
+
+    const words = currentSentence.trim().split(/\s+/);
+    const lastWord = words[words.length - 1] || '';
+
+    const baseWord = lastWord.slice(0, -1);
+
+    if (this.isCommonAbbreviation(baseWord)) {
+      return false;
+    }
+
+    return true;
+  }
+
+  private isCommonAbbreviation(word: string): boolean {
+    // Common titles
+    const titles = ['Dr', 'Mr', 'Mrs', 'Ms', 'Prof', 'Sr', 'Jr'];
+    if (titles.includes(word)) {
+      return true;
+    }
+
+    // Multi-character abbreviations with periods (U.S.A., a.m., p.m., etc.)
+    if (/^[A-Z](\.[A-Z])*$/.test(word) || /^[a-z](\.[a-z])*$/.test(word)) {
+      return true;
+    }
+
+    // Single capital letters (initials)
+    if (/^[A-Z]$/.test(word)) {
+      return true;
+    }
+
+    // Numbers (versions, decimals)
+    if (/^\d+$/.test(word)) {
+      return true;
+    }
+
+    // Time abbreviations
+    if (/^[ap]\.?m$/i.test(word)) {
+      return true;
+    }
+
+    return false;
+  }
+
+  /**
+   * Group sentences into chunks with integrated overlap processing
+   */
+  private groupSentencesIntoChunks(sentences: string[]): string[] {
+    const chunks: string[] = [];
+    let currentChunk: string[] = [];
+    let currentSize = 0;
+
+    const separator = ' ';
+
+    for (const sentence of sentences) {
+      const sentenceLength = this.lengthFunction(sentence);
+      const separatorLength = currentChunk.length > 0 ? this.lengthFunction(separator) : 0;
+      const totalLength = currentSize + sentenceLength + separatorLength;
+
+      // Handle oversized sentences with fallback strategies
+      if (sentenceLength > this.maxSize) {
+        if (currentChunk.length > 0) {
+          chunks.push(currentChunk.join(separator));
+          currentChunk = [];
+          currentSize = 0;
+        }
+
+        const fallbackChunks = this.handleOversizedSentence(sentence);
+        chunks.push(...fallbackChunks);
+        continue;
+      }
+
+      // If adding this sentence would exceed maxSize, finalize current chunk
+      if (currentChunk.length > 0 && totalLength > this.maxSize) {
+        chunks.push(currentChunk.join(separator));
+
+        const overlapSentences = this.calculateSentenceOverlap(currentChunk);
+        currentChunk = overlapSentences;
+        currentSize = this.calculateChunkSize(currentChunk);
+      }
+
+      currentChunk.push(sentence);
+      currentSize += sentenceLength + separatorLength;
+
+      // If we've reached our target size, consider finalizing the chunk
+      if (currentSize >= this.targetSize) {
+        chunks.push(currentChunk.join(separator));
+
+        const overlapSentences = this.calculateSentenceOverlap(currentChunk);
+        currentChunk = overlapSentences;
+        currentSize = this.calculateChunkSize(currentChunk);
+      }
+    }
+
+    if (currentChunk.length > 0) {
+      chunks.push(currentChunk.join(separator));
+    }
+
+    return chunks;
+  }
+
+  /**
+   * Handle oversized sentences with fallback strategies
+   */
+  private handleOversizedSentence(sentence: string): string[] {
+    // First fallback
+    if (this.fallbackToWords) {
+      const wordChunks = this.splitSentenceIntoWords(sentence);
+      if (wordChunks.length > 1) {
+        return wordChunks;
+      }
+    }
+
+    // Second fallback
+    if (this.fallbackToCharacters) {
+      return this.splitSentenceIntoCharacters(sentence);
+    }
+
+    // Last resort
+    console.warn(
+      `Sentence exceeds maxSize (${this.maxSize}) and fallbacks are disabled: "${sentence.substring(0, 50)}..."`,
+    );
+    return [sentence];
+  }
+
+  private splitSentenceIntoWords(sentence: string): string[] {
+    const words = sentence.split(/\s+/);
+    const chunks: string[] = [];
+    let currentChunk = '';
+
+    for (const word of words) {
+      const testChunk = currentChunk ? currentChunk + ' ' + word : word;
+
+      if (this.lengthFunction(testChunk) <= this.maxSize) {
+        currentChunk = testChunk;
+      } else {
+        if (currentChunk) {
+          chunks.push(currentChunk);
+        }
+
+        if (this.lengthFunction(word) > this.maxSize) {
+          if (this.fallbackToCharacters) {
+            chunks.push(...this.splitSentenceIntoCharacters(word));
+          } else {
+            chunks.push(word);
+          }
+          currentChunk = '';
+        } else {
+          currentChunk = word;
+        }
+      }
+    }
+
+    if (currentChunk) {
+      chunks.push(currentChunk);
+    }
+
+    return chunks;
+  }
+
+  private splitSentenceIntoCharacters(text: string): string[] {
+    const chunks: string[] = [];
+    let currentChunk = '';
+
+    for (const char of text) {
+      if (this.lengthFunction(currentChunk + char) <= this.maxSize) {
+        currentChunk += char;
+      } else {
+        if (currentChunk) {
+          chunks.push(currentChunk);
+        }
+        currentChunk = char;
+      }
+    }
+
+    if (currentChunk) {
+      chunks.push(currentChunk);
+    }
+
+    return chunks;
+  }
+
+  private calculateSentenceOverlap(currentChunk: string[]): string[] {
+    if (this.overlap === 0 || currentChunk.length === 0) {
+      return [];
+    }
+
+    const overlapSentences: string[] = [];
+    let overlapSize = 0;
+    const separator = ' ';
+
+    // Work backwards through sentences to build overlap
+    for (let i = currentChunk.length - 1; i >= 0; i--) {
+      const sentence = currentChunk[i];
+      if (!sentence) continue;
+
+      const sentenceLength = this.lengthFunction(sentence);
+      const separatorLength = overlapSentences.length > 0 ? this.lengthFunction(separator) : 0;
+
+      if (overlapSize + sentenceLength + separatorLength > this.overlap) {
+        break;
+      }
+
+      overlapSentences.unshift(sentence);
+      overlapSize += sentenceLength + separatorLength;
+    }
+
+    return overlapSentences;
+  }
+
+  private calculateChunkSize(sentences: string[]): number {
+    if (!sentences || sentences.length === 0) {
+      return 0;
+    }
+
+    let totalSize = 0;
+    const separator = ' ';
+
+    for (let i = 0; i < sentences.length; i++) {
+      const sentence = sentences[i]!;
+      totalSize += this.lengthFunction(sentence);
+
+      // Add separator length for all but the last sentence
+      if (i < sentences.length - 1) {
+        totalSize += this.lengthFunction(separator);
+      }
+    }
+
+    return totalSize;
+  }
+
+  splitText({ text }: { text: string }): string[] {
+    if (!text) return [];
+
+    const sentences = this.detectSentenceBoundaries(text);
+
+    const chunks = this.groupSentencesIntoChunks(sentences);
+
+    return chunks.filter(chunk => chunk.trim().length > 0);
+  }
+}

package/src/document/transformers/text.ts

@@ -1,11 +1,11 @@
 import { Document } from '../schema';
 
-import type { ChunkOptions } from '../types';
+import type { BaseChunkOptions } from '../types';
 
 import type { Transformer } from './transformer';
 
 export abstract class TextTransformer implements Transformer {
-  protected size: number;
+  protected maxSize: number;
   protected overlap: number;
   protected lengthFunction: (text: string) => number;
   protected keepSeparator: boolean | 'start' | 'end';
@@ -13,17 +13,17 @@ export abstract class TextTransformer implements Transformer {
   protected stripWhitespace: boolean;
 
   constructor({
-    size = 4000,
+    maxSize = 4000,
     overlap = 200,
     lengthFunction = (text: string) => text.length,
     keepSeparator = false,
     addStartIndex = false,
     stripWhitespace = true,
-  }: ChunkOptions) {
-    if (overlap > size) {
-      throw new Error(`Got a larger chunk overlap (${overlap}) than chunk size ` + `(${size}), should be smaller.`);
+  }: BaseChunkOptions) {
+    if (overlap > maxSize) {
+      throw new Error(`Got a larger chunk overlap (${overlap}) than chunk size ` + `(${maxSize}), should be smaller.`);
     }
-    this.size = size;
+    this.maxSize = maxSize;
     this.overlap = overlap;
     this.lengthFunction = lengthFunction;
     this.keepSeparator = keepSeparator;
@@ -104,9 +104,9 @@ export abstract class TextTransformer implements Transformer {
       const len = this.lengthFunction(d);
       const separatorLen = separator ? this.lengthFunction(separator) : 0;
 
-      if (total + len + (currentDoc.length > 0 ? separatorLen : 0) > this.size) {
-        if (total > this.size) {
-          console.warn(`Created a chunk of size ${total}, which is longer than the specified ${this.size}`);
+      if (total + len + (currentDoc.length > 0 ? separatorLen : 0) > this.maxSize) {
+        if (total > this.maxSize) {
+          console.warn(`Created a chunk of size ${total}, which is longer than the specified ${this.maxSize}`);
         }
 
         if (currentDoc.length > 0) {

package/src/document/transformers/token.ts

@@ -1,5 +1,6 @@
 import type { TiktokenModel, TiktokenEncoding, Tiktoken } from 'js-tiktoken';
 import { encodingForModel, getEncoding } from 'js-tiktoken';
+import type { TokenChunkOptions } from '../types';
 
 import { TextTransformer } from './text';
 
@@ -42,18 +43,11 @@ export class TokenTransformer extends TextTransformer {
     disallowedSpecial = 'all',
     options = {},
   }: {
-    encodingName: TiktokenEncoding;
+    encodingName?: TiktokenEncoding;
     modelName?: TiktokenModel;
     allowedSpecial?: Set<string> | 'all';
     disallowedSpecial?: Set<string> | 'all';
-    options: {
-      size?: number;
-      overlap?: number;
-      lengthFunction?: (text: string) => number;
-      keepSeparator?: boolean | 'start' | 'end';
-      addStartIndex?: boolean;
-      stripWhitespace?: boolean;
-    };
+    options: TokenChunkOptions;
   }) {
     super(options);
 
@@ -85,7 +79,7 @@ export class TokenTransformer extends TextTransformer {
 
     const tokenizer: Tokenizer = {
       overlap: this.overlap,
-      tokensPerChunk: this.size,
+      tokensPerChunk: this.maxSize,
       decode,
       encode,
     };
@@ -100,12 +94,7 @@ export class TokenTransformer extends TextTransformer {
   }: {
     encodingName?: TiktokenEncoding;
     modelName?: TiktokenModel;
-    options?: {
-      size?: number;
-      overlap?: number;
-      allowedSpecial?: Set<string> | 'all';
-      disallowedSpecial?: Set<string> | 'all';
-    };
+    options?: TokenChunkOptions;
   }): TokenTransformer {
     let tokenizer: Tiktoken;
 
@@ -139,7 +128,7 @@ export class TokenTransformer extends TextTransformer {
       allowedSpecial: options.allowedSpecial,
       disallowedSpecial: options.disallowedSpecial,
       options: {
-        size: options.size,
+        maxSize: options.maxSize,
         overlap: options.overlap,
         lengthFunction: tikTokenEncoder,
       },

package/src/document/types.ts

@@ -42,34 +42,85 @@ export type ExtractParams = {
   keywords?: KeywordExtractArgs | boolean;
 };
 
-export type ChunkOptions = {
-  headers?: [string, string][];
-  returnEachLine?: boolean;
-  sections?: [string, string][];
-  separator?: string;
-  separators?: string[];
-  isSeparatorRegex?: boolean;
+export type BaseChunkOptions = {
+  /**
+   * @deprecated Use `maxSize` instead. Will be removed in next major version.
+   */
   size?: number;
   maxSize?: number;
-  minSize?: number;
   overlap?: number;
   lengthFunction?: (text: string) => number;
   keepSeparator?: boolean | 'start' | 'end';
   addStartIndex?: boolean;
   stripWhitespace?: boolean;
+};
+
+export type CharacterChunkOptions = BaseChunkOptions & {
+  separator?: string;
+  isSeparatorRegex?: boolean;
+};
+
+export type RecursiveChunkOptions = BaseChunkOptions & {
+  separators?: string[];
+  isSeparatorRegex?: boolean;
   language?: Language;
-  ensureAscii?: boolean;
-  convertLists?: boolean;
+};
+
+export type TokenChunkOptions = BaseChunkOptions & {
   encodingName?: TiktokenEncoding;
   modelName?: TiktokenModel;
   allowedSpecial?: Set<string> | 'all';
   disallowedSpecial?: Set<string> | 'all';
+};
+
+export type MarkdownChunkOptions = BaseChunkOptions & {
+  headers?: [string, string][];
+  returnEachLine?: boolean;
   stripHeaders?: boolean;
 };
 
-export type ChunkStrategy = 'recursive' | 'character' | 'token' | 'markdown' | 'html' | 'json' | 'latex';
+export type HTMLChunkOptions = BaseChunkOptions &
+  (
+    | { headers: [string, string][]; sections?: never; returnEachLine?: boolean }
+    | { sections: [string, string][]; headers?: never }
+  ) & { returnEachLine?: boolean };
 
-export interface ChunkParams extends ChunkOptions {
-  strategy?: ChunkStrategy;
-  extract?: ExtractParams;
-}
+export type JsonChunkOptions = BaseChunkOptions & {
+  minSize?: number;
+  ensureAscii?: boolean;
+  convertLists?: boolean;
+};
+
+export type LatexChunkOptions = BaseChunkOptions & {};
+
+export type SentenceChunkOptions = BaseChunkOptions & {
+  maxSize: number; // Override to make required for sentence strategy
+  minSize?: number;
+  targetSize?: number;
+  sentenceEnders?: string[];
+  fallbackToWords?: boolean;
+  fallbackToCharacters?: boolean;
+};
+
+export type StrategyOptions = {
+  recursive: RecursiveChunkOptions;
+  character: CharacterChunkOptions;
+  token: TokenChunkOptions;
+  markdown: MarkdownChunkOptions;
+  html: HTMLChunkOptions;
+  json: JsonChunkOptions;
+  latex: LatexChunkOptions;
+  sentence: SentenceChunkOptions;
+};
+
+export type ChunkStrategy = 'recursive' | 'character' | 'token' | 'markdown' | 'html' | 'json' | 'latex' | 'sentence';
+
+export type ChunkParams =
+  | ({ strategy?: 'character' } & CharacterChunkOptions & { extract?: ExtractParams })
+  | ({ strategy: 'recursive' } & RecursiveChunkOptions & { extract?: ExtractParams })
+  | ({ strategy: 'token' } & TokenChunkOptions & { extract?: ExtractParams })
+  | ({ strategy: 'markdown' } & MarkdownChunkOptions & { extract?: ExtractParams })
+  | ({ strategy: 'html' } & HTMLChunkOptions & { extract?: ExtractParams })
+  | ({ strategy: 'json' } & JsonChunkOptions & { extract?: ExtractParams })
+  | ({ strategy: 'latex' } & LatexChunkOptions & { extract?: ExtractParams })
+  | ({ strategy: 'sentence' } & SentenceChunkOptions & { extract?: ExtractParams });

package/src/document/validation.ts

@@ -0,0 +1,147 @@
+import { z } from 'zod';
+import type { ChunkStrategy } from './types';
+
+function handleDeprecatedSize<T extends { size?: number; maxSize?: number }>(data: T): Omit<T, 'size'> {
+  if (data.size !== undefined) {
+    console.warn(
+      '[DEPRECATION] `size` is deprecated. Use `maxSize` instead. This will be removed in the next major version.',
+    );
+
+    if (data.maxSize === undefined) {
+      data.maxSize = data.size;
+    }
+  }
+
+  const { size, ...rest } = data;
+  return rest;
+}
+
+// Base options that apply to all strategies
+const baseChunkOptionsSchema = z.object({
+  size: z.number().positive().optional(),
+  maxSize: z.number().positive().optional(),
+  overlap: z.number().min(0).optional(),
+  lengthFunction: z.function().optional(),
+  keepSeparator: z.union([z.boolean(), z.literal('start'), z.literal('end')]).optional(),
+  addStartIndex: z.boolean().optional(),
+  stripWhitespace: z.boolean().optional(),
+});
+
+// Strategy-specific schemas
+const characterChunkOptionsSchema = baseChunkOptionsSchema
+  .extend({
+    separator: z.string().optional(),
+    isSeparatorRegex: z.boolean().optional(),
+  })
+  .strict();
+
+const recursiveChunkOptionsSchema = baseChunkOptionsSchema
+  .extend({
+    separators: z.array(z.string()).optional(),
+    isSeparatorRegex: z.boolean().optional(),
+    language: z.string().optional(),
+  })
+  .strict();
+
+const sentenceChunkOptionsSchema = baseChunkOptionsSchema
+  .extend({
+    maxSize: z.number().positive(),
+    minSize: z.number().positive().optional(),
+    targetSize: z.number().positive().optional(),
+    sentenceEnders: z.array(z.string()).optional(),
+    fallbackToWords: z.boolean().optional(),
+    fallbackToCharacters: z.boolean().optional(),
+  })
+  .strict();
+
+// Predicate to check for Set-like objects
+const isSetLike = (value: unknown): value is Set<any> => {
+  return (
+    typeof value === 'object' &&
+    value !== null &&
+    typeof (value as Set<any>).has === 'function' &&
+    typeof (value as Set<any>).add === 'function' &&
+    typeof (value as Set<any>).delete === 'function' &&
+    typeof (value as Set<any>).clear === 'function' &&
+    typeof (value as Set<any>).size === 'number'
+  );
+};
+
+// Zod schema for a Set or the literal 'all'
+const setOrAllSchema = z
+  .any()
+  .refine(value => value === 'all' || isSetLike(value), {
+    message: "Must be a Set object or the literal 'all'",
+  })
+  .optional();
+
+const tokenChunkOptionsSchema = baseChunkOptionsSchema
+  .extend({
+    encodingName: z.string().optional(),
+    modelName: z.string().optional(),
+    allowedSpecial: setOrAllSchema,
+    disallowedSpecial: setOrAllSchema,
+  })
+  .strict();
+
+const jsonChunkOptionsSchema = baseChunkOptionsSchema
+  .extend({
+    minSize: z.number().positive().optional(),
+    ensureAscii: z.boolean().optional(),
+    convertLists: z.boolean().optional(),
+  })
+  .strict();
+
+const htmlChunkOptionsSchema = baseChunkOptionsSchema
+  .extend({
+    headers: z.array(z.tuple([z.string(), z.string()])).optional(),
+    sections: z.array(z.tuple([z.string(), z.string()])).optional(),
+    returnEachLine: z.boolean().optional(),
+  })
+  .strict();
+
+const markdownChunkOptionsSchema = baseChunkOptionsSchema
+  .extend({
+    headers: z.array(z.tuple([z.string(), z.string()])).optional(),
+    returnEachLine: z.boolean().optional(),
+    stripHeaders: z.boolean().optional(),
+  })
+  .strict();
+
+const latexChunkOptionsSchema = baseChunkOptionsSchema.strict();
+
+// Strategy-specific validation schemas
+const validationSchemas = {
+  character: characterChunkOptionsSchema.transform(handleDeprecatedSize),
+  recursive: recursiveChunkOptionsSchema.transform(handleDeprecatedSize),
+  sentence: sentenceChunkOptionsSchema.transform(handleDeprecatedSize),
+  token: tokenChunkOptionsSchema.transform(handleDeprecatedSize),
+  json: jsonChunkOptionsSchema.transform(handleDeprecatedSize),
+  html: htmlChunkOptionsSchema.transform(handleDeprecatedSize),
+  markdown: markdownChunkOptionsSchema.transform(handleDeprecatedSize),
+  latex: latexChunkOptionsSchema.transform(handleDeprecatedSize),
+} as const;
+
+export function validateChunkParams(strategy: ChunkStrategy, params: any): void {
+  const schema = validationSchemas[strategy];
+  if (!schema) {
+    throw new Error(`Unknown chunking strategy: ${strategy}`);
+  }
+
+  const result = schema.safeParse(params);
+  if (!result.success) {
+    // Extract unrecognized keys for cleaner error message
+    const unrecognizedError = result.error.errors.find((e: any) => e.code === 'unrecognized_keys');
+    if (unrecognizedError && 'keys' in unrecognizedError) {
+      const keys = (unrecognizedError as any).keys.join(', ');
+      throw new Error(`Invalid parameters for ${strategy} strategy: '${keys}' not supported`);
+    }
+
+    // Fallback to general error message for other validation issues
+    const errorMessage = result.error.errors
+      .map((e: any) => `${e.path.length > 0 ? e.path.join('.') : 'parameter'}: ${e.message}`)
+      .join(', ');
+
+    throw new Error(`Invalid parameters for ${strategy} strategy: ${errorMessage}`);
+  }
+}