@mastra/rag 0.0.0-commonjs-20250227130920

package/src/document/transformers/markdown.ts

@@ -0,0 +1,244 @@
+import { Document } from 'llamaindex';
+
+import { Language } from '../types';
+
+import { RecursiveCharacterTransformer } from './character';
+
+interface LineType {
+  metadata: Record<string, string>;
+  content: string;
+}
+
+interface HeaderType {
+  level: number;
+  name: string;
+  data: string;
+}
+
+export class MarkdownTransformer extends RecursiveCharacterTransformer {
+  constructor(
+    options: {
+      chunkSize?: number;
+      chunkOverlap?: number;
+      lengthFunction?: (text: string) => number;
+      keepSeparator?: boolean | 'start' | 'end';
+      addStartIndex?: boolean;
+      stripWhitespace?: boolean;
+    } = {},
+  ) {
+    const separators = RecursiveCharacterTransformer.getSeparatorsForLanguage(Language.MARKDOWN);
+    super({ separators, isSeparatorRegex: true, options });
+  }
+}
+
+export class MarkdownHeaderTransformer {
+  private headersToSplitOn: [string, string][];
+  private returnEachLine: boolean;
+  private stripHeaders: boolean;
+
+  constructor(headersToSplitOn: [string, string][], returnEachLine: boolean = false, stripHeaders: boolean = true) {
+    this.headersToSplitOn = [...headersToSplitOn].sort((a, b) => b[0].length - a[0].length);
+    this.returnEachLine = returnEachLine;
+    this.stripHeaders = stripHeaders;
+  }
+
+  private aggregateLinesToChunks(lines: LineType[]): Document[] {
+    if (this.returnEachLine) {
+      return lines.flatMap(line => {
+        const contentLines = line.content.split('\n');
+        return contentLines
+          .filter(l => l.trim() !== '' || this.headersToSplitOn.some(([sep]) => l.trim().startsWith(sep)))
+          .map(
+            l =>
+              new Document({
+                text: l.trim(),
+                metadata: line.metadata,
+              }),
+          );
+      });
+    }
+
+    const aggregatedChunks: LineType[] = [];
+
+    for (const line of lines) {
+      if (
+        aggregatedChunks.length > 0 &&
+        JSON.stringify(aggregatedChunks?.[aggregatedChunks.length - 1]!.metadata) === JSON.stringify(line.metadata)
+      ) {
+        const aggChunk = aggregatedChunks[aggregatedChunks.length - 1];
+        aggChunk!.content += '  \n' + line.content;
+      } else if (
+        aggregatedChunks.length > 0 &&
+        JSON.stringify(aggregatedChunks?.[aggregatedChunks.length - 1]!.metadata) !== JSON.stringify(line.metadata) &&
+        Object.keys(aggregatedChunks?.[aggregatedChunks.length - 1]!.metadata).length <
+          Object.keys(line.metadata).length &&
+        aggregatedChunks?.[aggregatedChunks.length - 1]?.content?.split('\n')?.slice(-1)[0]![0] === '#' &&
+        !this.stripHeaders
+      ) {
+        if (aggregatedChunks && aggregatedChunks?.[aggregatedChunks.length - 1]) {
+          const aggChunk = aggregatedChunks[aggregatedChunks.length - 1];
+          if (aggChunk) {
+            aggChunk.content += '  \n' + line.content;
+            aggChunk.metadata = line.metadata;
+          }
+        }
+      } else {
+        aggregatedChunks.push(line);
+      }
+    }
+
+    return aggregatedChunks.map(
+      chunk =>
+        new Document({
+          text: chunk.content,
+          metadata: chunk.metadata,
+        }),
+    );
+  }
+
+  splitText({ text }: { text: string }): Document[] {
+    const lines = text.split('\n');
+    const linesWithMetadata: LineType[] = [];
+    let currentContent: string[] = [];
+    let currentMetadata: Record<string, string> = {};
+    const headerStack: HeaderType[] = [];
+    const initialMetadata: Record<string, string> = {};
+
+    let inCodeBlock = false;
+    let openingFence = '';
+
+    for (let i = 0; i < lines.length; i++) {
+      const line = lines[i]!;
+      const strippedLine = line.trim();
+
+      if (!inCodeBlock) {
+        if (
+          (strippedLine.startsWith('```') && strippedLine.split('```').length === 2) ||
+          strippedLine.startsWith('~~~')
+        ) {
+          inCodeBlock = true;
+          openingFence = strippedLine.startsWith('```') ? '```' : '~~~';
+        }
+      } else {
+        if (strippedLine.startsWith(openingFence)) {
+          inCodeBlock = false;
+          openingFence = '';
+        }
+      }
+
+      if (inCodeBlock) {
+        currentContent.push(line);
+        continue;
+      }
+
+      let headerMatched = false;
+      for (const [sep, name] of this.headersToSplitOn) {
+        if (strippedLine.startsWith(sep) && (strippedLine.length === sep.length || strippedLine[sep.length] === ' ')) {
+          headerMatched = true;
+
+          // If we have existing content, save it before processing the header
+          if (currentContent.length > 0) {
+            linesWithMetadata.push({
+              content: currentContent.join('\n'),
+              metadata: { ...currentMetadata },
+            });
+            currentContent = [];
+          }
+
+          if (name !== null) {
+            const currentHeaderLevel = (sep.match(/#/g) || []).length;
+
+            // Pop headers of lower or same level
+            while (headerStack.length > 0 && headerStack?.[headerStack.length - 1]!.level >= currentHeaderLevel) {
+              const poppedHeader = headerStack.pop()!;
+              if (poppedHeader.name in initialMetadata) {
+                delete initialMetadata[poppedHeader.name];
+              }
+            }
+
+            // Push current header
+            const header: HeaderType = {
+              level: currentHeaderLevel,
+              name,
+              data: strippedLine.slice(sep.length).trim(),
+            };
+            headerStack.push(header);
+            initialMetadata[name] = header.data;
+          }
+
+          // Always create a separate chunk for the header
+          linesWithMetadata.push({
+            content: line,
+            metadata: { ...currentMetadata, ...initialMetadata },
+          });
+
+          break;
+        }
+      }
+
+      if (!headerMatched) {
+        if (strippedLine || this.returnEachLine) {
+          currentContent.push(line);
+
+          if (this.returnEachLine) {
+            // In returnEachLine mode, flush each non-header line immediately
+            linesWithMetadata.push({
+              content: line,
+              metadata: { ...currentMetadata },
+            });
+            currentContent = [];
+          }
+        } else if (currentContent.length > 0) {
+          linesWithMetadata.push({
+            content: currentContent.join('\n'),
+            metadata: { ...currentMetadata },
+          });
+          currentContent = [];
+        }
+      }
+
+      currentMetadata = { ...initialMetadata };
+    }
+
+    // Handle any remaining content
+    if (currentContent.length > 0) {
+      linesWithMetadata.push({
+        content: currentContent.join('\n'),
+        metadata: currentMetadata,
+      });
+    }
+
+    return this.aggregateLinesToChunks(linesWithMetadata);
+  }
+
+  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], ...chunk.metadata };
+        documents.push(
+          new Document({
+            text: chunk.text,
+            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);
+  }
+}

package/src/document/transformers/text.ts

@@ -0,0 +1,134 @@
+import { Document } from 'llamaindex';
+
+import type { ChunkOptions } from '../types';
+
+import type { Transformer } from './transformer';
+
+export abstract class TextTransformer implements Transformer {
+  protected size: number;
+  protected overlap: number;
+  protected lengthFunction: (text: string) => number;
+  protected keepSeparator: boolean | 'start' | 'end';
+  protected addStartIndex: boolean;
+  protected stripWhitespace: boolean;
+
+  constructor({
+    size = 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.`);
+    }
+    this.size = size;
+    this.overlap = overlap;
+    this.lengthFunction = lengthFunction;
+    this.keepSeparator = keepSeparator;
+    this.addStartIndex = addStartIndex;
+    this.stripWhitespace = stripWhitespace;
+  }
+
+  setAddStartIndex(value: boolean): void {
+    this.addStartIndex = value;
+  }
+
+  abstract splitText({ text }: { text: string }): string[];
+
+  createDocuments(texts: string[], metadatas?: Record<string, any>[]): Document[] {
+    const _metadatas = metadatas || Array(texts.length).fill({});
+    const documents: Document[] = [];
+
+    texts.forEach((text, i) => {
+      let index = 0;
+      let previousChunkLen = 0;
+
+      this.splitText({ text }).forEach(chunk => {
+        const metadata = { ..._metadatas[i] };
+        if (this.addStartIndex) {
+          const offset = index + previousChunkLen - this.overlap;
+          index = text.indexOf(chunk, Math.max(0, offset));
+          metadata.startIndex = index;
+          previousChunkLen = chunk.length;
+        }
+        documents.push(
+          new Document({
+            text: chunk,
+            metadata,
+          }),
+        );
+      });
+    });
+
+    return documents;
+  }
+
+  splitDocuments(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);
+  }
+
+  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);
+  }
+
+  protected joinDocs(docs: string[], separator: string): string | null {
+    let text = docs.join(separator);
+    if (this.stripWhitespace) {
+      text = text.trim();
+    }
+    return text === '' ? null : text;
+  }
+
+  protected mergeSplits(splits: string[], separator: string): string[] {
+    const separatorLen = this.lengthFunction(separator);
+    const docs: string[] = [];
+    let currentDoc: string[] = [];
+    let total = 0;
+
+    splits.forEach(d => {
+      const len = this.lengthFunction(d);
+      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 (currentDoc.length > 0) {
+          const doc = this.joinDocs(currentDoc, separator);
+          if (doc !== null) {
+            docs.push(doc);
+          }
+          while (
+            total > this.overlap ||
+            (total + len + (currentDoc.length > 0 ? separatorLen : 0) > this.size && total > 0)
+          ) {
+            total -= this.lengthFunction(currentDoc?.[0]!) + (currentDoc.length > 1 ? separatorLen : 0);
+            currentDoc = currentDoc.slice(1);
+          }
+        }
+      }
+      currentDoc.push(d);
+      total += len + (currentDoc.length > 1 ? separatorLen : 0);
+    });
+
+    const doc = this.joinDocs(currentDoc, separator);
+    if (doc !== null) {
+      docs.push(doc);
+    }
+    return docs;
+  }
+}

package/src/document/transformers/token.ts

@@ -0,0 +1,148 @@
+import type { TiktokenModel, TiktokenEncoding, Tiktoken } from 'js-tiktoken';
+import { encodingForModel, getEncoding } from 'js-tiktoken';
+
+import { TextTransformer } from './text';
+
+interface Tokenizer {
+  overlap: number;
+  tokensPerChunk: number;
+  decode: (tokens: number[]) => string;
+  encode: (text: string) => number[];
+}
+
+export function splitTextOnTokens({ text, tokenizer }: { text: string; tokenizer: Tokenizer }): string[] {
+  const splits: string[] = [];
+  const inputIds = tokenizer.encode(text);
+  let startIdx = 0;
+  let curIdx = Math.min(startIdx + tokenizer.tokensPerChunk, inputIds.length);
+  let chunkIds = inputIds.slice(startIdx, curIdx);
+
+  while (startIdx < inputIds.length) {
+    splits.push(tokenizer.decode(chunkIds));
+    if (curIdx === inputIds.length) {
+      break;
+    }
+    startIdx += tokenizer.tokensPerChunk - tokenizer.overlap;
+    curIdx = Math.min(startIdx + tokenizer.tokensPerChunk, inputIds.length);
+    chunkIds = inputIds.slice(startIdx, curIdx);
+  }
+
+  return splits;
+}
+
+export class TokenTransformer extends TextTransformer {
+  private tokenizer: Tiktoken;
+  private allowedSpecial: Set<string> | 'all';
+  private disallowedSpecial: Set<string> | 'all';
+
+  constructor({
+    encodingName = 'cl100k_base',
+    modelName,
+    allowedSpecial = new Set(),
+    disallowedSpecial = 'all',
+    options = {},
+  }: {
+    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;
+    };
+  }) {
+    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`.');
+    }
+
+    this.allowedSpecial = allowedSpecial;
+    this.disallowedSpecial = disallowedSpecial;
+  }
+
+  splitText({ text }: { text: string }): string[] {
+    const encode = (text: string): number[] => {
+      const allowed = this.allowedSpecial === 'all' ? 'all' : Array.from(this.allowedSpecial);
+
+      const disallowed = this.disallowedSpecial === 'all' ? 'all' : Array.from(this.disallowedSpecial);
+
+      // If stripWhitespace is enabled, trim the text before encoding
+      const processedText = this.stripWhitespace ? text.trim() : text;
+      return Array.from(this.tokenizer.encode(processedText, allowed, disallowed));
+    };
+
+    const decode = (tokens: number[]): string => {
+      const text = this.tokenizer.decode(tokens);
+      return this.stripWhitespace ? text.trim() : text;
+    };
+
+    const tokenizer: Tokenizer = {
+      overlap: this.overlap,
+      tokensPerChunk: this.size,
+      decode,
+      encode,
+    };
+
+    return splitTextOnTokens({ text, tokenizer });
+  }
+
+  static fromTikToken({
+    encodingName = 'cl100k_base',
+    modelName,
+    options = {},
+  }: {
+    encodingName?: TiktokenEncoding;
+    modelName?: TiktokenModel;
+    options?: {
+      size?: number;
+      overlap?: number;
+      allowedSpecial?: Set<string> | 'all';
+      disallowedSpecial?: Set<string> | 'all';
+    };
+  }): TokenTransformer {
+    let tokenizer: Tiktoken;
+
+    try {
+      if (modelName) {
+        tokenizer = encodingForModel(modelName);
+      } else {
+        tokenizer = getEncoding(encodingName);
+      }
+    } catch {
+      throw new Error('Could not load tiktoken encoding. ' + 'Please install it with `npm install js-tiktoken`.');
+    }
+
+    const tikTokenEncoder = (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 TokenTransformer({
+      encodingName,
+      modelName,
+      allowedSpecial: options.allowedSpecial,
+      disallowedSpecial: options.disallowedSpecial,
+      options: {
+        size: options.size,
+        overlap: options.overlap,
+        lengthFunction: tikTokenEncoder,
+      },
+    });
+  }
+}

package/src/document/transformers/transformer.ts

@@ -0,0 +1,5 @@
+import type { Document } from 'llamaindex';
+
+export interface Transformer {
+  transformDocuments(documents: Document[]): Document[];
+}

package/src/document/types.ts

@@ -0,0 +1,103 @@
+import type { TiktokenEncoding, TiktokenModel } from 'js-tiktoken';
+import type {
+  LLM,
+  TitleCombinePrompt,
+  TitleExtractorPrompt,
+  SummaryPrompt,
+  QuestionExtractPrompt,
+  KeywordExtractPrompt,
+} from 'llamaindex';
+
+export enum Language {
+  CPP = 'cpp',
+  GO = 'go',
+  JAVA = 'java',
+  KOTLIN = 'kotlin',
+  JS = 'js',
+  TS = 'ts',
+  PHP = 'php',
+  PROTO = 'proto',
+  PYTHON = 'python',
+  RST = 'rst',
+  RUBY = 'ruby',
+  RUST = 'rust',
+  SCALA = 'scala',
+  SWIFT = 'swift',
+  MARKDOWN = 'markdown',
+  LATEX = 'latex',
+  HTML = 'html',
+  SOL = 'sol',
+  CSHARP = 'csharp',
+  COBOL = 'cobol',
+  C = 'c',
+  LUA = 'lua',
+  PERL = 'perl',
+  HASKELL = 'haskell',
+  ELIXIR = 'elixir',
+  POWERSHELL = 'powershell',
+}
+
+export type ExtractParams = {
+  title?: TitleExtractorsArgs | boolean;
+  summary?: SummaryExtractArgs | boolean;
+  questions?: QuestionAnswerExtractArgs | boolean;
+  keywords?: boolean | Record<string, any>;
+};
+
+export type ChunkOptions = {
+  headers?: [string, string][];
+  returnEachLine?: boolean;
+  sections?: [string, string][];
+  separator?: string;
+  separators?: string[];
+  isSeparatorRegex?: boolean;
+  size?: number;
+  maxSize?: number;
+  minSize?: number;
+  overlap?: number;
+  lengthFunction?: (text: string) => number;
+  keepSeparator?: boolean | 'start' | 'end';
+  addStartIndex?: boolean;
+  stripWhitespace?: boolean;
+  language?: Language;
+  ensureAscii?: boolean;
+  convertLists?: boolean;
+  encodingName?: TiktokenEncoding;
+  modelName?: TiktokenModel;
+  allowedSpecial?: Set<string> | 'all';
+  disallowedSpecial?: Set<string> | 'all';
+  stripHeaders?: boolean;
+};
+
+export type TitleExtractorsArgs = {
+  llm?: LLM;
+  nodes?: number;
+  nodeTemplate?: TitleExtractorPrompt['template'];
+  combineTemplate?: TitleCombinePrompt['template'];
+};
+
+export type SummaryExtractArgs = {
+  llm?: LLM;
+  summaries?: string[];
+  promptTemplate?: SummaryPrompt['template'];
+};
+
+export type QuestionAnswerExtractArgs = {
+  llm?: LLM;
+  questions?: number;
+  promptTemplate?: QuestionExtractPrompt['template'];
+  embeddingOnly?: boolean;
+};
+
+export type KeywordExtractArgs = {
+  llm?: LLM;
+  keywords?: number;
+  promptTemplate?: KeywordExtractPrompt['template'];
+};
+
+export type ChunkStrategy = 'recursive' | 'character' | 'token' | 'markdown' | 'html' | 'json' | 'latex';
+
+export interface ChunkParams extends ChunkOptions {
+  strategy?: ChunkStrategy;
+  extract?: ExtractParams;
+}