@mastra/rag 0.0.0-commonjs-20250227130920

package/src/document/document.ts

@@ -0,0 +1,281 @@
+import {
+  Document as Chunk,
+  IngestionPipeline,
+  KeywordExtractor,
+  QuestionsAnsweredExtractor,
+  SummaryExtractor,
+  TitleExtractor,
+} from 'llamaindex';
+
+import { CharacterTransformer, RecursiveCharacterTransformer } from './transformers/character';
+import { HTMLHeaderTransformer, HTMLSectionTransformer } from './transformers/html';
+import { RecursiveJsonTransformer } from './transformers/json';
+import { LatexTransformer } from './transformers/latex';
+import { MarkdownHeaderTransformer, MarkdownTransformer } from './transformers/markdown';
+import { TokenTransformer } from './transformers/token';
+import type { ChunkOptions, ChunkParams, ChunkStrategy, ExtractParams } from './types';
+
+export class MDocument {
+  private chunks: Chunk[];
+  private type: string; // e.g., 'text', 'html', 'markdown', 'json'
+
+  constructor({ docs, type }: { docs: { text: string; metadata?: Record<string, any> }[]; type: string }) {
+    this.chunks = docs.map(d => {
+      return new Chunk({ text: d.text, metadata: d.metadata });
+    });
+    this.type = type;
+  }
+
+  async extractMetadata({ title, summary, questions, keywords }: ExtractParams): Promise<MDocument> {
+    const transformations = [];
+
+    if (typeof summary !== 'undefined') {
+      transformations.push(new SummaryExtractor(typeof summary === 'boolean' ? {} : summary));
+    }
+
+    if (typeof questions !== 'undefined') {
+      transformations.push(new QuestionsAnsweredExtractor(typeof questions === 'boolean' ? {} : questions));
+    }
+
+    if (typeof keywords !== 'undefined') {
+      transformations.push(new KeywordExtractor(typeof keywords === 'boolean' ? {} : keywords));
+    }
+
+    if (typeof title !== 'undefined') {
+      transformations.push(new TitleExtractor(typeof title === 'boolean' ? {} : title));
+    }
+
+    const pipeline = new IngestionPipeline({
+      transformations,
+    });
+
+    const nodes = await pipeline.run({
+      documents: this.chunks,
+    });
+
+    this.chunks = this.chunks.map((doc, i) => {
+      return new Chunk({
+        text: doc.text,
+        metadata: {
+          ...doc.metadata,
+          ...(nodes?.[i]?.metadata || {}),
+        },
+      });
+    });
+
+    return this;
+  }
+
+  static fromText(text: string, metadata?: Record<string, any>): MDocument {
+    return new MDocument({
+      docs: [
+        {
+          text,
+          metadata,
+        },
+      ],
+      type: 'text',
+    });
+  }
+
+  static fromHTML(html: string, metadata?: Record<string, any>): MDocument {
+    return new MDocument({
+      docs: [
+        {
+          text: html,
+          metadata,
+        },
+      ],
+      type: 'html',
+    });
+  }
+
+  static fromMarkdown(markdown: string, metadata?: Record<string, any>): MDocument {
+    return new MDocument({
+      docs: [
+        {
+          text: markdown,
+          metadata,
+        },
+      ],
+      type: 'markdown',
+    });
+  }
+
+  static fromJSON(jsonString: string, metadata?: Record<string, any>): MDocument {
+    return new MDocument({
+      docs: [
+        {
+          text: jsonString,
+          metadata,
+        },
+      ],
+      type: 'json',
+    });
+  }
+
+  private defaultStrategy(): ChunkStrategy {
+    switch (this.type) {
+      case 'html':
+        return 'html';
+      case 'markdown':
+        return 'markdown';
+      case 'json':
+        return 'json';
+      case 'latex':
+        return 'latex';
+      default:
+        return 'recursive';
+    }
+  }
+
+  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}`);
+    }
+  }
+
+  async chunkRecursive(options?: ChunkOptions): Promise<void> {
+    if (options?.language) {
+      const rt = RecursiveCharacterTransformer.fromLanguage(options.language, options);
+      const textSplit = rt.transformDocuments(this.chunks);
+      this.chunks = textSplit;
+      return;
+    }
+
+    const rt = new RecursiveCharacterTransformer({
+      separators: options?.separators,
+      isSeparatorRegex: options?.isSeparatorRegex,
+      options,
+    });
+    const textSplit = rt.transformDocuments(this.chunks);
+    this.chunks = textSplit;
+  }
+
+  async chunkCharacter(options?: ChunkOptions): Promise<void> {
+    const rt = new CharacterTransformer({
+      separator: options?.separator,
+      isSeparatorRegex: options?.isSeparatorRegex,
+      options,
+    });
+    const textSplit = rt.transformDocuments(this.chunks);
+    this.chunks = textSplit;
+  }
+
+  async chunkHTML(options?: ChunkOptions): Promise<void> {
+    if (options?.headers?.length) {
+      const rt = new HTMLHeaderTransformer(options.headers, options?.returnEachLine);
+
+      const textSplit = rt.transformDocuments(this.chunks);
+      this.chunks = textSplit;
+      return;
+    }
+
+    if (options?.sections?.length) {
+      const rt = new HTMLSectionTransformer(options.sections);
+
+      const textSplit = rt.transformDocuments(this.chunks);
+      this.chunks = textSplit;
+      return;
+    }
+
+    throw new Error('HTML chunking requires either headers or sections to be specified');
+  }
+
+  async chunkJSON(options?: ChunkOptions): Promise<void> {
+    if (!options?.maxSize) {
+      throw new Error('JSON chunking requires maxSize to be specified');
+    }
+
+    const rt = new RecursiveJsonTransformer({
+      maxSize: options?.maxSize,
+      minSize: options?.minSize,
+    });
+
+    const textSplit = rt.transformDocuments({
+      documents: this.chunks,
+      ensureAscii: options?.ensureAscii,
+      convertLists: options?.convertLists,
+    });
+
+    this.chunks = textSplit;
+  }
+
+  async chunkLatex(options?: ChunkOptions): Promise<void> {
+    const rt = new LatexTransformer(options);
+    const textSplit = rt.transformDocuments(this.chunks);
+    this.chunks = textSplit;
+  }
+
+  async chunkToken(options?: ChunkOptions): Promise<void> {
+    const rt = TokenTransformer.fromTikToken({
+      options,
+      encodingName: options?.encodingName,
+      modelName: options?.modelName,
+    });
+    const textSplit = rt.transformDocuments(this.chunks);
+    this.chunks = textSplit;
+  }
+
+  async chunkMarkdown(options?: ChunkOptions): Promise<void> {
+    if (options?.headers) {
+      const rt = new MarkdownHeaderTransformer(options.headers, options?.returnEachLine, options?.stripHeaders);
+      const textSplit = rt.transformDocuments(this.chunks);
+      this.chunks = textSplit;
+      return;
+    }
+
+    const rt = new MarkdownTransformer(options);
+    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();
+
+    // Apply the appropriate chunking strategy
+    await this.chunkBy(strategy, chunkOptions);
+
+    if (extract) {
+      await this.extractMetadata(extract);
+    }
+
+    return this.chunks;
+  }
+
+  getDocs(): Chunk[] {
+    return this.chunks;
+  }
+
+  getText(): string[] {
+    return this.chunks.map(doc => doc.text);
+  }
+
+  getMetadata(): Record<string, any>[] {
+    return this.chunks.map(doc => doc.metadata);
+  }
+}

package/src/document/index.ts

@@ -0,0 +1,2 @@
+export * from './document';
+export * from './types';

package/src/document/transformers/character.ts

@@ -0,0 +1,279 @@
+import { Language } from '../types';
+import type { ChunkOptions } from '../types';
+
+import { TextTransformer } from './text';
+
+function splitTextWithRegex(text: string, separator: string, keepSeparator: boolean | 'start' | 'end'): string[] {
+  if (!separator) {
+    return text.split('');
+  }
+
+  if (!keepSeparator) {
+    return text.split(new RegExp(separator)).filter(s => s !== '');
+  }
+
+  if (!text) {
+    return [];
+  }
+
+  // Split with capturing group to keep separators
+  const splits = text.split(new RegExp(`(${separator})`));
+  const result: string[] = [];
+
+  if (keepSeparator === 'end') {
+    // Process all complete pairs
+    for (let i = 0; i < splits.length - 1; i += 2) {
+      if (i + 1 < splits.length) {
+        // Current text + separator
+        const chunk = splits[i] + (splits[i + 1] || '');
+        if (chunk) result.push(chunk);
+      }
+    }
+    // Handle the last element if it exists and isn't a separator
+    if (splits.length % 2 === 1 && splits[splits.length - 1]) {
+      result.push(splits?.[splits.length - 1]!);
+    }
+  } else {
+    if (splits[0]) result.push(splits[0]);
+
+    for (let i = 1; i < splits.length - 1; i += 2) {
+      const separator = splits[i];
+      const text = splits[i + 1];
+      if (separator && text) {
+        result.push(separator + text);
+      }
+    }
+  }
+
+  return result.filter(s => s !== '');
+}
+
+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);
+    this.separator = separator;
+    this.isSeparatorRegex = isSeparatorRegex;
+  }
+
+  splitText({ text }: { text: string }): string[] {
+    // First, split the text into initial chunks
+    const separator = this.isSeparatorRegex ? this.separator : this.separator.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+
+    const initialSplits = splitTextWithRegex(text, separator, this.keepSeparator);
+
+    // 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) {
+        chunks.push(split);
+      } else {
+        // If a single split is too large, split it further
+        const subChunks = this.__splitChunk(split);
+        chunks.push(...subChunks);
+      }
+    }
+
+    return chunks;
+  }
+
+  private __splitChunk(text: string): string[] {
+    const chunks: string[] = [];
+    let currentChunk = '';
+
+    // Split by characters if no other separator is suitable
+    const chars = text.split('');
+
+    for (const char of chars) {
+      if (this.lengthFunction(currentChunk + char) <= this.size) {
+        currentChunk += char;
+      } else {
+        if (currentChunk) {
+          chunks.push(currentChunk);
+        }
+        currentChunk = char;
+      }
+    }
+
+    if (currentChunk) {
+      chunks.push(currentChunk);
+    }
+
+    return chunks;
+  }
+}
+
+export class RecursiveCharacterTransformer extends TextTransformer {
+  protected separators: string[];
+  protected isSeparatorRegex: boolean;
+
+  constructor({
+    separators,
+    isSeparatorRegex = false,
+    options = {},
+  }: {
+    separators?: string[];
+    isSeparatorRegex?: boolean;
+    options?: ChunkOptions;
+  }) {
+    super(options);
+    this.separators = separators || ['\n\n', '\n', ' ', ''];
+    this.isSeparatorRegex = isSeparatorRegex;
+  }
+
+  private _splitText(text: string, separators: string[]): string[] {
+    const finalChunks: string[] = [];
+
+    let separator = separators?.[separators.length - 1]!;
+    let newSeparators: string[] = [];
+
+    for (let i = 0; i < separators.length; i++) {
+      const s = separators[i]!;
+      const _separator = this.isSeparatorRegex ? s : s?.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+
+      if (s === '') {
+        separator = s;
+        break;
+      }
+
+      if (new RegExp(_separator).test(text)) {
+        separator = s;
+        newSeparators = separators.slice(i + 1);
+        break;
+      }
+    }
+
+    const _separator = this.isSeparatorRegex ? separator : separator?.replace(/[.*+?^${}()|[\]\\]/g, '\\$&');
+
+    const splits = splitTextWithRegex(text, _separator, this.keepSeparator);
+
+    const goodSplits: string[] = [];
+    const mergeSeparator = this.keepSeparator ? '' : separator;
+
+    for (const s of splits) {
+      if (this.lengthFunction(s) < this.size) {
+        goodSplits.push(s);
+      } else {
+        if (goodSplits.length > 0) {
+          const mergedText = this.mergeSplits(goodSplits, mergeSeparator);
+          finalChunks.push(...mergedText);
+          goodSplits.length = 0;
+        }
+        if (newSeparators.length === 0) {
+          finalChunks.push(s);
+        } else {
+          const otherInfo = this._splitText(s, newSeparators);
+          finalChunks.push(...otherInfo);
+        }
+      }
+    }
+
+    if (goodSplits.length > 0) {
+      const mergedText = this.mergeSplits(goodSplits, mergeSeparator);
+      finalChunks.push(...mergedText);
+    }
+
+    return finalChunks;
+  }
+
+  splitText({ text }: { text: string }): string[] {
+    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 {
+    const separators = RecursiveCharacterTransformer.getSeparatorsForLanguage(language);
+    return new RecursiveCharacterTransformer({ separators, isSeparatorRegex: true, options });
+  }
+
+  static getSeparatorsForLanguage(language: Language): string[] {
+    switch (language) {
+      case Language.MARKDOWN:
+        return [
+          // First, try to split along Markdown headings (starting with level 2)
+          '\n#{1,6} ',
+          // End of code block
+          '```\n',
+          // Horizontal lines
+          '\n\\*\\*\\*+\n',
+          '\n---+\n',
+          '\n___+\n',
+          // Note that this splitter doesn't handle horizontal lines defined
+          // by *three or more* of ***, ---, or ___, but this is not handled
+          '\n\n',
+          '\n',
+          ' ',
+          '',
+        ];
+      case Language.CPP:
+      case Language.C:
+        return [
+          '\nclass ',
+          '\nvoid ',
+          '\nint ',
+          '\nfloat ',
+          '\ndouble ',
+          '\nif ',
+          '\nfor ',
+          '\nwhile ',
+          '\nswitch ',
+          '\ncase ',
+          '\n\n',
+          '\n',
+          ' ',
+          '',
+        ];
+      case Language.TS:
+        return [
+          '\nenum ',
+          '\ninterface ',
+          '\nnamespace ',
+          '\ntype ',
+          '\nclass ',
+          '\nfunction ',
+          '\nconst ',
+          '\nlet ',
+          '\nvar ',
+          '\nif ',
+          '\nfor ',
+          '\nwhile ',
+          '\nswitch ',
+          '\ncase ',
+          '\ndefault ',
+          '\n\n',
+          '\n',
+          ' ',
+          '',
+        ];
+      // ... (add other language cases following the same pattern)
+      default:
+        throw new Error(`Language ${language} is not supported! Please choose from ${Object.values(Language)}`);
+    }
+  }
+}