@mastra/rag 1.2.2 → 1.2.3-alpha.1

package/src/document/document.ts

@@ -1,335 +0,0 @@
-import { TitleExtractor, SummaryExtractor, QuestionsAnsweredExtractor, KeywordExtractor } from './extractors';
-import type { BaseNode } from './schema';
-import { Document as Chunk, NodeRelationship, ObjectType } from './schema';
-
-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 { SemanticMarkdownTransformer } from './transformers/semantic-markdown';
-import { SentenceTransformer } from './transformers/sentence';
-import { TokenTransformer } from './transformers/token';
-import type {
-  ChunkParams,
-  ChunkStrategy,
-  ExtractParams,
-  HTMLChunkOptions,
-  RecursiveChunkOptions,
-  CharacterChunkOptions,
-  TokenChunkOptions,
-  MarkdownChunkOptions,
-  SemanticMarkdownChunkOptions,
-  JsonChunkOptions,
-  LatexChunkOptions,
-  SentenceChunkOptions,
-  StrategyOptions,
-} from './types';
-import { validateChunkParams } from './validation';
-
-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));
-      this.chunks = this.chunks.map(doc =>
-        doc?.metadata?.docId
-          ? new Chunk({
-              ...doc,
-              relationships: {
-                [NodeRelationship.SOURCE]: {
-                  nodeId: doc.metadata.docId,
-                  nodeType: ObjectType.DOCUMENT,
-                  metadata: doc.metadata,
-                },
-              },
-            })
-          : doc,
-      );
-    }
-
-    let nodes: BaseNode[] = this.chunks;
-    for (const extractor of transformations) {
-      nodes = await extractor.processNodes(nodes);
-    }
-
-    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 _strategyMap?: { [S in ChunkStrategy]: (options?: StrategyOptions[S]) => Promise<void> };
-
-  private get strategyMap() {
-    if (!this._strategyMap) {
-      this._strategyMap = {
-        recursive: options => this.chunkRecursive(options),
-        character: options => this.chunkCharacter(options),
-        token: options => this.chunkToken(options),
-        markdown: options => this.chunkMarkdown(options),
-        html: options => this.chunkHTML(options),
-        json: options => this.chunkJSON(options),
-        latex: options => this.chunkLatex(options),
-        sentence: options => this.chunkSentence(options),
-        'semantic-markdown': options => this.chunkSemanticMarkdown(options),
-      };
-    }
-    return this._strategyMap;
-  }
-
-  private async chunkBy<K extends ChunkStrategy>(strategy: K, options?: StrategyOptions[K]): Promise<void> {
-    const chunkingFunc = this.strategyMap[strategy];
-    if (chunkingFunc) {
-      await chunkingFunc(options);
-    } else {
-      throw new Error(`Unknown strategy: ${strategy}`);
-    }
-  }
-
-  async chunkRecursive(options?: RecursiveChunkOptions): 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(options);
-    const textSplit = rt.transformDocuments(this.chunks);
-    this.chunks = textSplit;
-  }
-
-  async chunkCharacter(options?: CharacterChunkOptions): Promise<void> {
-    const rt = new CharacterTransformer({
-      ...options,
-      separator: options?.separator,
-      isSeparatorRegex: options?.isSeparatorRegex,
-    });
-    const textSplit = rt.transformDocuments(this.chunks);
-    this.chunks = textSplit;
-  }
-
-  async chunkHTML(options?: HTMLChunkOptions): Promise<void> {
-    if (options?.headers?.length) {
-      const rt = new HTMLHeaderTransformer(options as HTMLChunkOptions & { headers: [string, string][] });
-
-      const textSplit = rt.transformDocuments(this.chunks);
-      this.chunks = textSplit;
-      return;
-    }
-
-    if (options?.sections?.length) {
-      const rt = new HTMLSectionTransformer(options as HTMLChunkOptions & { sections: [string, string][] });
-
-      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?: JsonChunkOptions): 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?: LatexChunkOptions): Promise<void> {
-    const rt = new LatexTransformer(options);
-    const textSplit = rt.transformDocuments(this.chunks);
-    this.chunks = textSplit;
-  }
-
-  async chunkToken(options?: TokenChunkOptions): 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?: MarkdownChunkOptions): 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 chunkSentence(options?: SentenceChunkOptions): Promise<void> {
-    if (!options?.maxSize) {
-      throw new Error('Sentence chunking requires maxSize to be specified');
-    }
-
-    const rt = new SentenceTransformer({
-      minSize: options?.minSize,
-      maxSize: options?.maxSize,
-      targetSize: options?.targetSize,
-      overlap: options?.overlap,
-      sentenceEnders: options?.sentenceEnders,
-      fallbackToWords: options?.fallbackToWords,
-      fallbackToCharacters: options?.fallbackToCharacters,
-      keepSeparator: options?.keepSeparator,
-      lengthFunction: options?.lengthFunction,
-      addStartIndex: options?.addStartIndex,
-      stripWhitespace: options?.stripWhitespace,
-    });
-
-    const textSplit = rt.transformDocuments(this.chunks);
-    this.chunks = textSplit;
-  }
-
-  async chunkSemanticMarkdown(options?: SemanticMarkdownChunkOptions): Promise<void> {
-    const rt = SemanticMarkdownTransformer.fromTikToken({
-      options,
-      encodingName: options?.encodingName,
-      modelName: options?.modelName,
-    });
-    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();
-
-    validateChunkParams(strategy, chunkOptions);
-
-    // 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/extractors/base.ts

@@ -1,30 +0,0 @@
-import type { BaseNode } from '../schema';
-
-/*
- * Abstract class for all extractors.
- */
-export abstract class BaseExtractor {
-  isTextNodeOnly: boolean = true;
-
-  abstract extract(nodes: BaseNode[]): Promise<Record<string, any>[]>;
-
-  /**
-   *
-   * @param nodes Nodes to extract metadata from.
-   * @returns Metadata extracted from the nodes.
-   */
-  async processNodes(nodes: BaseNode[]): Promise<BaseNode[]> {
-    let newNodes: BaseNode[] = nodes;
-
-    const curMetadataList = await this.extract(newNodes);
-
-    for (const idx in newNodes) {
-      newNodes[idx]!.metadata = {
-        ...newNodes[idx]!.metadata,
-        ...curMetadataList[idx],
-      };
-    }
-
-    return newNodes;
-  }
-}

package/src/document/extractors/index.ts

@@ -1,5 +0,0 @@
-export { TitleExtractor } from './title';
-export { SummaryExtractor } from './summary';
-export { QuestionsAnsweredExtractor } from './questions';
-export { KeywordExtractor } from './keywords';
-export type { KeywordExtractArgs, QuestionAnswerExtractArgs, SummaryExtractArgs, TitleExtractorsArgs } from './types';

package/src/document/extractors/keywords.test.ts

@@ -1,125 +0,0 @@
-import { createOpenAI } from '@ai-sdk/openai';
-import { describe, it, expect, vi } from 'vitest';
-import { TextNode } from '../schema';
-import { KeywordExtractor } from './keywords';
-
-const openai = createOpenAI({
-  apiKey: process.env.OPENAI_API_KEY,
-});
-
-const model = openai('gpt-4o');
-
-vi.setConfig({ testTimeout: 50_000, hookTimeout: 50_000 });
-
-describe('KeywordExtractor', () => {
-  it('can use a custom model for keywords extraction', async () => {
-    const extractor = new KeywordExtractor({ llm: model });
-    const node = new TextNode({ text: 'The quick brown fox jumps over the lazy dog.' });
-    const result = await extractor.extractKeywordsFromNodes(node);
-    expect(result).toHaveProperty('excerptKeywords');
-    expect(result.excerptKeywords.length).toBeGreaterThan(0);
-  });
-
-  it('handles empty input gracefully', async () => {
-    const extractor = new KeywordExtractor();
-    const node = new TextNode({ text: '' });
-    const result = await extractor.extractKeywordsFromNodes(node);
-    expect(result.excerptKeywords).toBe('');
-  });
-
-  it('supports prompt customization', async () => {
-    const extractor = new KeywordExtractor({
-      promptTemplate: 'List keywords in: {context}. Limit to {maxKeywords}.',
-    });
-    const node = new TextNode({ text: 'Test document for prompt customization.' });
-    const result = await extractor.extractKeywordsFromNodes(node);
-    expect(result).toHaveProperty('excerptKeywords');
-    expect(typeof result.excerptKeywords).toBe('string');
-    expect(result.excerptKeywords.length).toBeGreaterThan(0);
-  });
-
-  it('extracts keywords from text', async () => {
-    const extractor = new KeywordExtractor();
-    const node = new TextNode({ text: 'The quick brown fox jumps over the lazy dog.' });
-    const result = await extractor.extractKeywordsFromNodes(node);
-    expect(result).toHaveProperty('excerptKeywords');
-    expect(typeof result.excerptKeywords).toBe('string');
-    expect(result.excerptKeywords.length).toBeGreaterThan(0);
-  });
-  it(
-    'handles very long input',
-    {
-      timeout: 60_000,
-    },
-    async () => {
-      const extractor = new KeywordExtractor();
-      const longText = 'A'.repeat(1000);
-      const node = new TextNode({ text: longText });
-      const result = await extractor.extractKeywordsFromNodes(node);
-      expect(result).toHaveProperty('excerptKeywords');
-      expect(typeof result.excerptKeywords).toBe('string');
-      expect(result.excerptKeywords.length).toBeGreaterThan(0);
-    },
-  );
-
-  it('handles whitespace only input', async () => {
-    const extractor = new KeywordExtractor();
-    const node = new TextNode({ text: '    ' });
-    const result = await extractor.extractKeywordsFromNodes(node);
-    expect(result.excerptKeywords).toBe('');
-  });
-
-  it('handles special characters and emojis', async () => {
-    const extractor = new KeywordExtractor();
-    const node = new TextNode({ text: '🚀✨🔥' });
-    const result = await extractor.extractKeywordsFromNodes(node);
-    expect(result).toHaveProperty('excerptKeywords');
-    expect(typeof result.excerptKeywords).toBe('string');
-    expect(result.excerptKeywords.length).toBeGreaterThan(0);
-  });
-
-  it('handles numbers only', async () => {
-    const extractor = new KeywordExtractor();
-    const node = new TextNode({ text: '1234567890' });
-    const result = await extractor.extractKeywordsFromNodes(node);
-    expect(result).toHaveProperty('excerptKeywords');
-    expect(typeof result.excerptKeywords).toBe('string');
-    expect(result.excerptKeywords.length).toBeGreaterThan(0);
-  });
-
-  it('handles HTML tags', async () => {
-    const extractor = new KeywordExtractor();
-    const node = new TextNode({ text: '<h1>Test</h1>' });
-    const result = await extractor.extractKeywordsFromNodes(node);
-    expect(result).toHaveProperty('excerptKeywords');
-    expect(typeof result.excerptKeywords).toBe('string');
-    expect(result.excerptKeywords.length).toBeGreaterThan(0);
-  });
-
-  it('handles non-English text', async () => {
-    const extractor = new KeywordExtractor();
-    const node = new TextNode({ text: '这是一个测试文档。' });
-    const result = await extractor.extractKeywordsFromNodes(node);
-    expect(result).toHaveProperty('excerptKeywords');
-    expect(typeof result.excerptKeywords).toBe('string');
-    expect(result.excerptKeywords.length).toBeGreaterThan(0);
-  });
-
-  it('handles duplicate/repeated text', async () => {
-    const extractor = new KeywordExtractor();
-    const node = new TextNode({ text: 'repeat repeat repeat' });
-    const result = await extractor.extractKeywordsFromNodes(node);
-    expect(result).toHaveProperty('excerptKeywords');
-    expect(typeof result.excerptKeywords).toBe('string');
-    expect(result.excerptKeywords.length).toBeGreaterThan(0);
-  });
-
-  it('handles only punctuation', async () => {
-    const extractor = new KeywordExtractor();
-    const node = new TextNode({ text: '!!!???...' });
-    const result = await extractor.extractKeywordsFromNodes(node);
-    expect(result).toHaveProperty('excerptKeywords');
-    expect(typeof result.excerptKeywords).toBe('string');
-    expect(result.excerptKeywords.length).toBeGreaterThan(0);
-  });
-});

package/src/document/extractors/keywords.ts

@@ -1,126 +0,0 @@
-import { Agent } from '@mastra/core/agent';
-import type { MastraLanguageModel } from '@mastra/core/agent';
-import { defaultKeywordExtractPrompt, PromptTemplate } from '../prompts';
-import type { KeywordExtractPrompt } from '../prompts';
-import type { BaseNode } from '../schema';
-import { TextNode } from '../schema';
-import { BaseExtractor } from './base';
-import { baseLLM } from './types';
-import type { KeywordExtractArgs } from './types';
-
-type ExtractKeyword = {
-  /**
-   * Comma-separated keywords extracted from the node. May be empty if extraction fails.
-   */
-  excerptKeywords: string;
-};
-
-/**
- * Extract keywords from a list of nodes.
- */
-export class KeywordExtractor extends BaseExtractor {
-  llm: MastraLanguageModel;
-  keywords: number = 5;
-  promptTemplate: KeywordExtractPrompt;
-
-  /**
-   * Constructor for the KeywordExtractor class.
-   * @param {MastraLanguageModel} llm MastraLanguageModel instance.
-   * @param {number} keywords Number of keywords to extract.
-   * @param {string} [promptTemplate] Optional custom prompt template (must include {context})
-   * @throws {Error} If keywords is less than 1.
-   */
-  constructor(options?: KeywordExtractArgs) {
-    if (options?.keywords && options.keywords < 1) throw new Error('Keywords must be greater than 0');
-
-    super();
-
-    this.llm = options?.llm ?? baseLLM;
-    this.keywords = options?.keywords ?? 5;
-    this.promptTemplate = options?.promptTemplate
-      ? new PromptTemplate({
-          templateVars: ['context', 'maxKeywords'],
-          template: options.promptTemplate,
-        })
-      : defaultKeywordExtractPrompt;
-  }
-
-  /**
-   *
-   * @param node Node to extract keywords from.
-   * @returns Keywords extracted from the node.
-   */
-  /**
-   * Extract keywords from a node. Returns an object with a comma-separated string of keywords, or an empty string if extraction fails.
-   * Adds error handling for malformed/empty LLM output.
-   */
-  async extractKeywordsFromNodes(node: BaseNode): Promise<ExtractKeyword> {
-    const text = node.getContent();
-    if (!text || text.trim() === '') {
-      return { excerptKeywords: '' };
-    }
-    if (this.isTextNodeOnly && !(node instanceof TextNode)) {
-      return { excerptKeywords: '' };
-    }
-
-    let keywords = '';
-    try {
-      const miniAgent = new Agent({
-        model: this.llm,
-        name: 'keyword-extractor',
-        instructions:
-          'You are a keyword extractor. You are given a node and you need to extract the keywords from the node.',
-      });
-
-      if (this.llm.specificationVersion === 'v2') {
-        const result = await miniAgent.generateVNext(
-          [
-            {
-              role: 'user',
-              content: this.promptTemplate.format({
-                context: node.getContent(),
-                maxKeywords: this.keywords.toString(),
-              }),
-            },
-          ],
-          { format: 'mastra' },
-        );
-        keywords = result.text;
-      } else {
-        const result = await miniAgent.generate([
-          {
-            role: 'user',
-            content: this.promptTemplate.format({ context: node.getContent(), maxKeywords: this.keywords.toString() }),
-          },
-        ]);
-        keywords = result.text;
-      }
-
-      if (!keywords) {
-        console.warn('Keyword extraction LLM output returned empty');
-        return { excerptKeywords: '' };
-      }
-
-      return { excerptKeywords: keywords.trim() };
-    } catch (err) {
-      console.warn('Keyword extraction failed:', err);
-      return { excerptKeywords: '' };
-    }
-  }
-
-  /**
-   *
-   * @param nodes Nodes to extract keywords from.
-   * @returns Keywords extracted from the nodes.
-   */
-  /**
-   * Extract keywords from an array of nodes. Always returns an array (may be empty).
-   * @param nodes Nodes to extract keywords from.
-   * @returns Array of keyword extraction results.
-   */
-  async extract(nodes: BaseNode[]): Promise<Array<ExtractKeyword>> {
-    if (!Array.isArray(nodes) || nodes.length === 0) return [];
-    const results = await Promise.all(nodes.map(node => this.extractKeywordsFromNodes(node)));
-    return results;
-  }
-}