@mastra/rag 1.2.3-alpha.0 → 1.2.3-alpha.2

package/src/document/transformers/markdown.ts

@@ -1,239 +0,0 @@
-import { Document } from '../schema';
-
-import { Language } from '../types';
-import type { BaseChunkOptions } 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: BaseChunkOptions = {}) {
-    const separators = RecursiveCharacterTransformer.getSeparatorsForLanguage(Language.MARKDOWN);
-    super({ ...options, separators, isSeparatorRegex: true });
-  }
-}
-
-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) {
-      const lastLine = aggregatedChunks[aggregatedChunks.length - 1]?.content?.split('\n')?.slice(-1)[0]?.trim();
-      const lastChunkIsHeader = lastLine ? this.headersToSplitOn.some(([sep]) => lastLine.startsWith(sep)) : false;
-      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 &&
-        lastChunkIsHeader
-      ) {
-        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;
-          }
-
-          // Only add header to linesWithMetadata if stripHeaders is false
-          if (!this.stripHeaders) {
-            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 = [];
-        }
-      }
-
-      // Reset metadata for next line
-      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/semantic-markdown.ts

@@ -1,227 +0,0 @@
-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,
-    });
-  }
-}