@mastra/rag 0.0.0-commonjs-20250227130920

package/src/document/transformers/html.ts

@@ -0,0 +1,346 @@
+import { Document } from 'llamaindex';
+import { parse } from 'node-html-better-parser';
+
+import { RecursiveCharacterTransformer } from './character';
+
+interface ElementType {
+  url: string;
+  xpath: string;
+  content: string;
+  metadata: Record<string, string>;
+}
+
+export class HTMLHeaderTransformer {
+  private headersToSplitOn: [string, string][];
+  private returnEachElement: boolean;
+
+  constructor(headersToSplitOn: [string, string][], returnEachElement: boolean = false) {
+    this.returnEachElement = returnEachElement;
+    this.headersToSplitOn = [...headersToSplitOn].sort();
+  }
+
+  splitText({ text }: { text: string }): Document[] {
+    const root = parse(text);
+
+    const headerFilter = this.headersToSplitOn.map(([header]) => header);
+    const headerMapping = Object.fromEntries(this.headersToSplitOn);
+
+    const elements: ElementType[] = [];
+    const headers = root.querySelectorAll(headerFilter.join(','));
+
+    headers.forEach(header => {
+      let content = '';
+      const parentNode = header.parentNode;
+
+      if (parentNode && parentNode.childNodes) {
+        let foundHeader = false;
+        for (const node of parentNode.childNodes) {
+          // Start collecting content after we find our header
+          if (node === header) {
+            foundHeader = true;
+            continue;
+          }
+
+          // If we found our header and hit another header, stop
+          // @ts-expect-error - node.tagName is not defined on type Node
+          if (foundHeader && node.tagName && headerFilter.includes(node.tagName.toLowerCase())) {
+            break;
+          }
+
+          // Collect content between headers
+          if (foundHeader) {
+            content += this.getTextContent(node) + ' ';
+          }
+        }
+      }
+
+      elements.push({
+        url: text,
+        xpath: this.getXPath(header),
+        content: content.trim(),
+        metadata: {
+          [headerMapping?.[header.tagName.toLowerCase()]!]: header.text || '',
+        },
+      });
+    });
+
+    return this.returnEachElement
+      ? elements.map(
+          el =>
+            new Document({
+              text: el.content,
+              metadata: { ...el.metadata, xpath: el.xpath },
+            }),
+        )
+      : this.aggregateElementsToChunks(elements);
+  }
+
+  private getXPath(element: any): string {
+    if (!element) return '';
+
+    const parts: string[] = [];
+    let current = element;
+
+    while (current && current.tagName) {
+      let index = 1;
+      const parent = current.parentNode;
+
+      if (parent && parent.childNodes) {
+        // Count preceding siblings with same tag
+        for (const sibling of parent.childNodes) {
+          if (sibling === current) break;
+          if (sibling.tagName === current.tagName) {
+            index++;
+          }
+        }
+      }
+
+      parts.unshift(`${current.tagName.toLowerCase()}[${index}]`);
+      current = current.parentNode;
+    }
+
+    return '/' + parts.join('/');
+  }
+
+  private getTextContent(element: any): string {
+    if (!element) return '';
+
+    // For text nodes, return their content
+    if (!element.tagName) {
+      return element.text || '';
+    }
+
+    // For element nodes, combine their text with children's text
+    let content = element.text || '';
+
+    if (element.childNodes) {
+      for (const child of element.childNodes) {
+        const childText = this.getTextContent(child);
+        if (childText) {
+          content += ' ' + childText;
+        }
+      }
+    }
+
+    return content.trim();
+  }
+
+  private aggregateElementsToChunks(elements: ElementType[]): Document[] {
+    const aggregatedChunks: ElementType[] = [];
+
+    for (const element of elements) {
+      if (
+        aggregatedChunks.length > 0 &&
+        JSON.stringify(aggregatedChunks[aggregatedChunks.length - 1]!.metadata) === JSON.stringify(element.metadata)
+      ) {
+        // If the last element has the same metadata, append content
+        aggregatedChunks[aggregatedChunks.length - 1]!.content += '  \n' + element.content;
+      } else {
+        // Otherwise, add as new element
+        aggregatedChunks.push({ ...element });
+      }
+    }
+
+    return aggregatedChunks.map(
+      chunk =>
+        new Document({
+          text: chunk.content,
+          metadata: { ...chunk.metadata, xpath: chunk.xpath },
+        }),
+    );
+  }
+
+  createDocuments(texts: string[], metadatas?: Record<string, any>[]): Document[] {
+    const _metadatas = metadatas || Array(texts.length).fill({});
+    const documents: Document[] = [];
+
+    for (let i = 0; i < texts.length; i++) {
+      const chunks = this.splitText({ text: texts[i]! });
+      for (const chunk of chunks) {
+        const metadata = { ...(_metadatas[i] || {}) };
+        const chunkMetadata = chunk.metadata;
+
+        if (chunkMetadata) {
+          for (const [key, value] of Object.entries(chunkMetadata || {})) {
+            if (value === '#TITLE#') {
+              chunkMetadata[key] = metadata['Title'];
+            }
+          }
+        }
+
+        documents.push(
+          new Document({
+            text: chunk.text!,
+            metadata: { ...metadata, ...chunkMetadata },
+          }),
+        );
+      }
+    }
+
+    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);
+  }
+}
+
+export class HTMLSectionTransformer {
+  private headersToSplitOn: Record<string, string>;
+  private options: Record<string, any>;
+
+  constructor(headersToSplitOn: [string, string][], options: Record<string, any> = {}) {
+    this.headersToSplitOn = Object.fromEntries(headersToSplitOn.map(([tag, name]) => [tag.toLowerCase(), name]));
+    this.options = options;
+  }
+
+  splitText(text: string): Document[] {
+    const sections = this.splitHtmlByHeaders(text);
+
+    return sections.map(
+      section =>
+        new Document({
+          text: section.content,
+          metadata: {
+            [this.headersToSplitOn[section.tagName.toLowerCase()]!]: section.header,
+            xpath: section.xpath,
+          },
+        }),
+    );
+  }
+
+  private getXPath(element: any): string {
+    const parts: string[] = [];
+    let current = element;
+
+    while (current && current.nodeType === 1) {
+      let index = 1;
+      let sibling = current.previousSibling;
+
+      while (sibling) {
+        if (sibling.nodeType === 1 && sibling.tagName === current.tagName) {
+          index++;
+        }
+        sibling = sibling.previousSibling;
+      }
+
+      if (current.tagName) {
+        parts.unshift(`${current.tagName.toLowerCase()}[${index}]`);
+      }
+      current = current.parentNode;
+    }
+
+    return '/' + parts.join('/');
+  }
+
+  private splitHtmlByHeaders(htmlDoc: string): Array<{
+    header: string;
+    content: string;
+    tagName: string;
+    xpath: string;
+  }> {
+    const sections: Array<{
+      header: string;
+      content: string;
+      tagName: string;
+      xpath: string;
+    }> = [];
+
+    const root = parse(htmlDoc);
+    const headers = Object.keys(this.headersToSplitOn);
+    const headerElements = root.querySelectorAll(headers.join(','));
+
+    headerElements.forEach((headerElement, index) => {
+      const header = headerElement.text?.trim() || '';
+      const tagName = headerElement.tagName;
+      const xpath = this.getXPath(headerElement);
+      let content = '';
+
+      // @ts-expect-error - nextElementSibling is not defined on type Element
+      let currentElement = headerElement.nextElementSibling;
+      const nextHeader = headerElements[index + 1];
+
+      while (currentElement && (!nextHeader || currentElement !== nextHeader)) {
+        if (currentElement.text) {
+          content += currentElement.text.trim() + ' ';
+        }
+        currentElement = currentElement.nextElementSibling;
+      }
+
+      content = content.trim();
+      sections.push({
+        header,
+        content,
+        tagName,
+        xpath,
+      });
+    });
+
+    return sections;
+  }
+
+  async splitDocuments(documents: Document[]): Promise<Document[]> {
+    const texts: string[] = [];
+    const metadatas: Record<string, any>[] = [];
+
+    for (const doc of documents) {
+      texts.push(doc.text);
+      metadatas.push(doc.metadata);
+    }
+    const results = await this.createDocuments(texts, metadatas);
+    const textSplitter = new RecursiveCharacterTransformer({ options: this.options });
+
+    return textSplitter.splitDocuments(results);
+  }
+
+  createDocuments(texts: string[], metadatas?: Record<string, any>[]): Document[] {
+    const _metadatas = metadatas || Array(texts.length).fill({});
+    const documents: Document[] = [];
+
+    for (let i = 0; i < texts.length; i++) {
+      const chunks = this.splitText(texts[i]!);
+      for (const chunk of chunks) {
+        const metadata = { ...(_metadatas[i] || {}) };
+
+        const chunkMetadata = chunk.metadata;
+
+        if (chunkMetadata) {
+          for (const [key, value] of Object.entries(chunkMetadata || {})) {
+            if (value === '#TITLE#') {
+              chunkMetadata[key] = metadata['Title'];
+            }
+          }
+        }
+
+        documents.push(
+          new Document({
+            text: chunk.text!,
+            metadata: { ...metadata, ...chunkMetadata },
+          }),
+        );
+      }
+    }
+
+    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/json.ts

@@ -0,0 +1,265 @@
+import { Document } from 'llamaindex';
+
+export class RecursiveJsonTransformer {
+  private maxSize: number;
+  private minSize: number;
+
+  constructor({ maxSize = 2000, minSize }: { maxSize: number; minSize?: number }) {
+    this.maxSize = maxSize;
+    this.minSize = minSize ?? Math.max(maxSize - 200, 50);
+  }
+
+  private static jsonSize(data: Record<string, any>): number {
+    const seen = new WeakSet();
+
+    function getStringifiableData(obj: any): any {
+      if (obj === null || typeof obj !== 'object') {
+        return obj;
+      }
+
+      if (seen.has(obj)) {
+        return '[Circular]';
+      }
+
+      seen.add(obj);
+
+      if (Array.isArray(obj)) {
+        const safeArray = [];
+        for (const item of obj) {
+          safeArray.push(getStringifiableData(item));
+        }
+        return safeArray;
+      }
+
+      const safeObj: Record<string, any> = {};
+      for (const key in obj) {
+        if (Object.prototype.hasOwnProperty.call(obj, key)) {
+          safeObj[key] = getStringifiableData(obj[key]);
+        }
+      }
+      return safeObj;
+    }
+
+    const stringifiable = getStringifiableData(data);
+    return JSON.stringify(stringifiable).length;
+  }
+
+  /**
+   * Transform JSON data while handling circular references
+   */
+  public transform(data: Record<string, any>): Record<string, any> {
+    const size = RecursiveJsonTransformer.jsonSize(data);
+
+    const seen = new WeakSet();
+
+    function createSafeCopy(obj: any): any {
+      if (obj === null || typeof obj !== 'object') {
+        return obj;
+      }
+
+      if (seen.has(obj)) {
+        return '[Circular]';
+      }
+
+      seen.add(obj);
+
+      if (Array.isArray(obj)) {
+        return obj.map(item => createSafeCopy(item));
+      }
+
+      const copy: Record<string, any> = {};
+      for (const key in obj) {
+        if (Object.prototype.hasOwnProperty.call(obj, key)) {
+          copy[key] = createSafeCopy(obj[key]);
+        }
+      }
+      return copy;
+    }
+
+    return {
+      size,
+      data: createSafeCopy(data),
+    };
+  }
+
+  /**
+   * Set a value in a nested dictionary based on the given path
+   */
+  private static setNestedDict(d: Record<string, any>, path: string[], value: any): void {
+    let current = d;
+    for (const key of path.slice(0, -1)) {
+      current[key] = current[key] || {};
+      current = current[key];
+    }
+    current[path[path.length - 1]!] = value;
+  }
+
+  /**
+   * Convert lists in the JSON structure to dictionaries with index-based keys
+   */
+  private listToDictPreprocessing(data: any): any {
+    if (data && typeof data === 'object') {
+      if (Array.isArray(data)) {
+        return Object.fromEntries(data.map((item, index) => [String(index), this.listToDictPreprocessing(item)]));
+      }
+      return Object.fromEntries(Object.entries(data).map(([k, v]) => [k, this.listToDictPreprocessing(v)]));
+    }
+    return data;
+  }
+
+  /**
+   * Split json into maximum size dictionaries while preserving structure
+   */
+  private jsonSplit({
+    data,
+    currentPath = [],
+    chunks = [{}],
+  }: {
+    data: Record<string, any>;
+    currentPath?: string[];
+    chunks?: Record<string, any>[];
+  }): Record<string, any>[] {
+    if (data && typeof data === 'object' && !Array.isArray(data)) {
+      for (const [key, value] of Object.entries(data)) {
+        const newPath = [...currentPath, key];
+        const chunkSize = RecursiveJsonTransformer.jsonSize(chunks[chunks.length - 1] || {});
+        const size = RecursiveJsonTransformer.jsonSize({ [key]: value });
+        const remaining = this.maxSize - chunkSize;
+
+        if (size < remaining) {
+          // Add item to current chunk
+          RecursiveJsonTransformer.setNestedDict(chunks[chunks.length - 1] || {}, newPath, value);
+        } else {
+          if (chunkSize >= this.minSize) {
+            // Chunk is big enough, start a new chunk
+            chunks.push({});
+          }
+          // Iterate
+          this.jsonSplit({
+            data: typeof value === 'object' ? value : { [key]: value },
+            currentPath: newPath,
+            chunks,
+          });
+        }
+      }
+    } else {
+      RecursiveJsonTransformer.setNestedDict(chunks[chunks.length - 1] || {}, currentPath, data);
+    }
+    return chunks;
+  }
+
+  /**
+   * Splits JSON into a list of JSON chunks
+   */
+  splitJson({
+    jsonData,
+    convertLists = false,
+  }: {
+    jsonData: Record<string, any>;
+    convertLists?: boolean;
+  }): Record<string, any>[] {
+    const processedData = convertLists ? this.listToDictPreprocessing(jsonData) : jsonData;
+
+    const chunks = this.jsonSplit({ data: processedData });
+
+    if (Object.keys(chunks[chunks.length - 1] || {}).length === 0) {
+      chunks.pop();
+    }
+
+    return chunks;
+  }
+
+  private escapeNonAscii(obj: any): any {
+    if (typeof obj === 'string') {
+      return obj.replace(/[\u0080-\uffff]/g, char => {
+        return `\\u${char.charCodeAt(0).toString(16).padStart(4, '0')}`;
+      });
+    }
+
+    if (typeof obj === 'object' && obj !== null) {
+      return Object.fromEntries(Object.entries(obj).map(([key, value]) => [key, this.escapeNonAscii(value)]));
+    }
+
+    return obj;
+  }
+  /**
+   * Splits JSON into a list of JSON formatted strings
+   */
+  splitText({
+    jsonData,
+    convertLists = false,
+    ensureAscii = true,
+  }: {
+    jsonData: Record<string, any>;
+    convertLists?: boolean;
+    ensureAscii?: boolean;
+  }): string[] {
+    const chunks = this.splitJson({ jsonData, convertLists });
+
+    if (ensureAscii) {
+      const escapedChunks = chunks.map(chunk => this.escapeNonAscii(chunk));
+      return escapedChunks.map(chunk => JSON.stringify(chunk));
+    }
+
+    return chunks.map(chunk => JSON.stringify(chunk));
+  }
+
+  /**
+   * Create documents from a list of json objects
+   */
+  createDocuments({
+    texts,
+    convertLists = false,
+    ensureAscii = true,
+    metadatas,
+  }: {
+    texts: string[];
+    convertLists?: boolean;
+    ensureAscii?: boolean;
+    metadatas?: Record<string, any>[];
+  }): Document[] {
+    const _metadatas = metadatas || Array(texts.length).fill({});
+    const documents: Document[] = [];
+
+    texts.forEach((text, i) => {
+      const chunks = this.splitText({ jsonData: JSON.parse(text), convertLists, ensureAscii });
+      chunks.forEach(chunk => {
+        const metadata = { ...(_metadatas[i] || {}) };
+        documents.push(
+          new Document({
+            text: chunk,
+            metadata,
+          }),
+        );
+      });
+    });
+
+    return documents;
+  }
+
+  transformDocuments({
+    ensureAscii,
+    documents,
+    convertLists,
+  }: {
+    ensureAscii?: boolean;
+    convertLists?: boolean;
+    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,
+
+      ensureAscii,
+      convertLists,
+    });
+  }
+}

package/src/document/transformers/latex.ts

@@ -0,0 +1,19 @@
+import { Language } 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;
+    } = {},
+  ) {
+    const separators = RecursiveCharacterTransformer.getSeparatorsForLanguage(Language.LATEX);
+    super({ separators, isSeparatorRegex: true, options });
+  }
+}