@mastra/rag 1.2.2 → 1.2.3-alpha.1

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

@@ -1,120 +0,0 @@
-import { createOpenAI } from '@ai-sdk/openai';
-import { describe, it, expect, vi } from 'vitest';
-import { TextNode } from '../schema';
-import { QuestionsAnsweredExtractor } from './questions';
-
-const openai = createOpenAI({
-  apiKey: process.env.OPENAI_API_KEY,
-});
-
-const model = openai('gpt-4o');
-
-vi.setConfig({ testTimeout: 100_000, hookTimeout: 100_000 });
-
-describe('QuestionsAnsweredExtractor', () => {
-  it('can use a custom model for questions extraction', async () => {
-    const extractor = new QuestionsAnsweredExtractor({ llm: model });
-    const node = new TextNode({ text: 'What is the capital of Spain?' });
-    const result = await extractor.extractQuestionsFromNode(node);
-    expect(result).toHaveProperty('questionsThisExcerptCanAnswer');
-    expect(result.questionsThisExcerptCanAnswer.length).toBeGreaterThan(0);
-  });
-
-  it('extracts questions', async () => {
-    const extractor = new QuestionsAnsweredExtractor();
-    const node = new TextNode({ text: 'What is the capital of France? What is the color of the sky?' });
-    const result = await extractor.extractQuestionsFromNode(node);
-    expect(result).toHaveProperty('questionsThisExcerptCanAnswer');
-    expect(typeof result.questionsThisExcerptCanAnswer).toBe('string');
-    expect(result.questionsThisExcerptCanAnswer.length).toBeGreaterThan(0);
-  });
-
-  it('handles empty input gracefully', async () => {
-    const extractor = new QuestionsAnsweredExtractor();
-    const node = new TextNode({ text: '' });
-    const result = await extractor.extractQuestionsFromNode(node);
-    expect(result).toHaveProperty('questionsThisExcerptCanAnswer');
-    expect(result.questionsThisExcerptCanAnswer).toBe('');
-  });
-
-  it('supports prompt customization', async () => {
-    const extractor = new QuestionsAnsweredExtractor({
-      promptTemplate: 'List questions in: {context}. Limit to {numQuestions}.',
-    });
-    const node = new TextNode({ text: 'Test document for prompt customization.' });
-    const result = await extractor.extractQuestionsFromNode(node);
-    expect(result).toHaveProperty('questionsThisExcerptCanAnswer');
-    expect(typeof result.questionsThisExcerptCanAnswer).toBe('string');
-    expect(result.questionsThisExcerptCanAnswer.length).toBeGreaterThan(0);
-  });
-  it('handles very long input', async () => {
-    const extractor = new QuestionsAnsweredExtractor();
-    const longText = 'A'.repeat(1000);
-    const node = new TextNode({ text: longText });
-    const result = await extractor.extractQuestionsFromNode(node);
-    expect(result).toHaveProperty('questionsThisExcerptCanAnswer');
-    expect(typeof result.questionsThisExcerptCanAnswer).toBe('string');
-    expect(result.questionsThisExcerptCanAnswer.length).toBeGreaterThan(0);
-  });
-
-  it('handles whitespace only input', async () => {
-    const extractor = new QuestionsAnsweredExtractor();
-    const node = new TextNode({ text: '    ' });
-    const result = await extractor.extractQuestionsFromNode(node);
-    expect(result.questionsThisExcerptCanAnswer).toBe('');
-  });
-
-  it('handles special characters and emojis', async () => {
-    const extractor = new QuestionsAnsweredExtractor();
-    const node = new TextNode({ text: '🚀✨🔥' });
-    const result = await extractor.extractQuestionsFromNode(node);
-    expect(result).toHaveProperty('questionsThisExcerptCanAnswer');
-    expect(typeof result.questionsThisExcerptCanAnswer).toBe('string');
-    expect(result.questionsThisExcerptCanAnswer.length).toBeGreaterThan(0);
-  });
-
-  it('handles numbers only', async () => {
-    const extractor = new QuestionsAnsweredExtractor();
-    const node = new TextNode({ text: '1234567890' });
-    const result = await extractor.extractQuestionsFromNode(node);
-    expect(result).toHaveProperty('questionsThisExcerptCanAnswer');
-    expect(typeof result.questionsThisExcerptCanAnswer).toBe('string');
-    expect(result.questionsThisExcerptCanAnswer.length).toBeGreaterThan(0);
-  });
-
-  it('handles HTML tags', async () => {
-    const extractor = new QuestionsAnsweredExtractor();
-    const node = new TextNode({ text: '<h1>Test</h1>' });
-    const result = await extractor.extractQuestionsFromNode(node);
-    expect(result).toHaveProperty('questionsThisExcerptCanAnswer');
-    expect(typeof result.questionsThisExcerptCanAnswer).toBe('string');
-    expect(result.questionsThisExcerptCanAnswer.length).toBeGreaterThan(0);
-  });
-
-  it('handles non-English text', async () => {
-    const extractor = new QuestionsAnsweredExtractor();
-    const node = new TextNode({ text: '这是一个测试文档。' });
-    const result = await extractor.extractQuestionsFromNode(node);
-    expect(result).toHaveProperty('questionsThisExcerptCanAnswer');
-    expect(typeof result.questionsThisExcerptCanAnswer).toBe('string');
-    expect(result.questionsThisExcerptCanAnswer.length).toBeGreaterThan(0);
-  });
-
-  it('handles duplicate/repeated text', async () => {
-    const extractor = new QuestionsAnsweredExtractor();
-    const node = new TextNode({ text: 'repeat repeat repeat' });
-    const result = await extractor.extractQuestionsFromNode(node);
-    expect(result).toHaveProperty('questionsThisExcerptCanAnswer');
-    expect(typeof result.questionsThisExcerptCanAnswer).toBe('string');
-    expect(result.questionsThisExcerptCanAnswer.length).toBeGreaterThan(0);
-  });
-
-  it('handles only punctuation', async () => {
-    const extractor = new QuestionsAnsweredExtractor();
-    const node = new TextNode({ text: '!!!???...' });
-    const result = await extractor.extractQuestionsFromNode(node);
-    expect(result).toHaveProperty('questionsThisExcerptCanAnswer');
-    expect(typeof result.questionsThisExcerptCanAnswer).toBe('string');
-    expect(result.questionsThisExcerptCanAnswer.length).toBeGreaterThan(0);
-  });
-});

package/src/document/extractors/questions.ts

@@ -1,111 +0,0 @@
-import { Agent } from '@mastra/core/agent';
-import type { MastraLanguageModel } from '@mastra/core/agent';
-import { PromptTemplate, defaultQuestionExtractPrompt } from '../prompts';
-import type { QuestionExtractPrompt } from '../prompts';
-import type { BaseNode } from '../schema';
-import { TextNode } from '../schema';
-import { BaseExtractor } from './base';
-import { baseLLM, STRIP_REGEX } from './types';
-import type { QuestionAnswerExtractArgs } from './types';
-
-type ExtractQuestion = {
-  /**
-   * Questions extracted from the node as a string (may be empty if extraction fails).
-   */
-  questionsThisExcerptCanAnswer: string;
-};
-
-/**
- * Extract questions from a list of nodes.
- */
-export class QuestionsAnsweredExtractor extends BaseExtractor {
-  llm: MastraLanguageModel;
-  questions: number = 5;
-  promptTemplate: QuestionExtractPrompt;
-  embeddingOnly: boolean = false;
-
-  /**
-   * Constructor for the QuestionsAnsweredExtractor class.
-   * @param {MastraLanguageModel} llm MastraLanguageModel instance.
-   * @param {number} questions Number of questions to generate.
-   * @param {QuestionExtractPrompt['template']} promptTemplate Optional custom prompt template (should include {context}).
-   * @param {boolean} embeddingOnly Whether to use metadata for embeddings only.
-   */
-  constructor(options?: QuestionAnswerExtractArgs) {
-    if (options?.questions && options.questions < 1) throw new Error('Questions must be greater than 0');
-
-    super();
-
-    this.llm = options?.llm ?? baseLLM;
-    this.questions = options?.questions ?? 5;
-    this.promptTemplate = options?.promptTemplate
-      ? new PromptTemplate({
-          templateVars: ['numQuestions', 'context'],
-          template: options.promptTemplate,
-        }).partialFormat({
-          numQuestions: '5',
-        })
-      : defaultQuestionExtractPrompt;
-    this.embeddingOnly = options?.embeddingOnly ?? false;
-  }
-
-  /**
-   * Extract answered questions from a node.
-   * @param {BaseNode} node Node to extract questions from.
-   * @returns {Promise<Array<ExtractQuestion> | Array<{}>>} Questions extracted from the node.
-   */
-  async extractQuestionsFromNode(node: BaseNode): Promise<ExtractQuestion> {
-    const text = node.getContent();
-    if (!text || text.trim() === '') {
-      return { questionsThisExcerptCanAnswer: '' };
-    }
-    if (this.isTextNodeOnly && !(node instanceof TextNode)) {
-      return { questionsThisExcerptCanAnswer: '' };
-    }
-
-    const contextStr = node.getContent();
-
-    const prompt = this.promptTemplate.format({
-      context: contextStr,
-      numQuestions: this.questions.toString(),
-    });
-
-    const miniAgent = new Agent({
-      model: this.llm,
-      name: 'question-extractor',
-      instructions:
-        'You are a question extractor. You are given a node and you need to extract the questions from the node.',
-    });
-
-    let questionsText = '';
-    if (this.llm.specificationVersion === 'v2') {
-      const result = await miniAgent.generateVNext([{ role: 'user', content: prompt }], { format: 'mastra' });
-      questionsText = result.text;
-    } else {
-      const result = await miniAgent.generate([{ role: 'user', content: prompt }]);
-      questionsText = result.text;
-    }
-
-    if (!questionsText) {
-      console.warn('Question extraction LLM output returned empty');
-      return { questionsThisExcerptCanAnswer: '' };
-    }
-
-    const result = questionsText.replace(STRIP_REGEX, '').trim();
-
-    return {
-      questionsThisExcerptCanAnswer: result,
-    };
-  }
-
-  /**
-   * Extract answered questions from a list of nodes.
-   * @param {BaseNode[]} nodes Nodes to extract questions from.
-   * @returns {Promise<Array<ExtractQuestion> | Array<{}>>} Questions extracted from the nodes.
-   */
-  async extract(nodes: BaseNode[]): Promise<Array<ExtractQuestion> | Array<object>> {
-    const results = await Promise.all(nodes.map(node => this.extractQuestionsFromNode(node)));
-
-    return results;
-  }
-}

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

@@ -1,107 +0,0 @@
-import { createOpenAI } from '@ai-sdk/openai';
-import { describe, it, expect, vi } from 'vitest';
-import { TextNode } from '../schema';
-import { SummaryExtractor } from './summary';
-
-const openai = createOpenAI({
-  apiKey: process.env.OPENAI_API_KEY,
-});
-
-const model = openai('gpt-4o');
-
-vi.setConfig({ testTimeout: 10_000, hookTimeout: 10_000 });
-
-describe('SummaryExtractor', () => {
-  it('can use a custom model from the test suite', async () => {
-    const extractor = new SummaryExtractor({ llm: model });
-    const node = new TextNode({ text: 'A summary test using a custom model.' });
-    const summary = await extractor.generateNodeSummary(node);
-    expect(typeof summary).toBe('string');
-    expect(summary.length).toBeGreaterThan(0);
-  });
-  it('extracts summary from normal text', async () => {
-    const extractor = new SummaryExtractor();
-    const node = new TextNode({ text: 'This is a test document.' });
-    const summary = await extractor.generateNodeSummary(node);
-    expect(typeof summary).toBe('string');
-    expect(summary.length).toBeGreaterThan(0);
-  });
-
-  it('handles empty input gracefully', async () => {
-    const extractor = new SummaryExtractor();
-    const node = new TextNode({ text: '' });
-    const summary = await extractor.generateNodeSummary(node);
-    expect(summary).toBe('');
-  });
-
-  it('supports prompt customization', async () => {
-    const extractor = new SummaryExtractor({ promptTemplate: 'Summarize: {context}' });
-    const node = new TextNode({ text: 'Test document for prompt customization.' });
-    const summary = await extractor.generateNodeSummary(node);
-    expect(typeof summary).toBe('string');
-    expect(summary.length).toBeGreaterThan(0);
-  });
-
-  it('handles very long input', async () => {
-    const extractor = new SummaryExtractor();
-    const longText = 'A'.repeat(1000);
-    const node = new TextNode({ text: longText });
-    const summary = await extractor.generateNodeSummary(node);
-    expect(typeof summary).toBe('string');
-  });
-
-  it('handles whitespace only input', async () => {
-    const extractor = new SummaryExtractor();
-    const node = new TextNode({ text: '    ' });
-    const summary = await extractor.generateNodeSummary(node);
-    expect(summary).toBe('');
-  });
-
-  it('handles special characters and emojis', async () => {
-    const extractor = new SummaryExtractor();
-    const node = new TextNode({ text: '🚀✨🔥' });
-    const summary = await extractor.generateNodeSummary(node);
-    expect(typeof summary).toBe('string');
-    expect(summary.length).toBeGreaterThan(0);
-  });
-
-  it('handles numbers only', async () => {
-    const extractor = new SummaryExtractor();
-    const node = new TextNode({ text: '1234567890' });
-    const summary = await extractor.generateNodeSummary(node);
-    expect(typeof summary).toBe('string');
-    expect(summary.length).toBeGreaterThan(0);
-  });
-
-  it('handles HTML tags', async () => {
-    const extractor = new SummaryExtractor();
-    const node = new TextNode({ text: '<h1>Test</h1>' });
-    const summary = await extractor.generateNodeSummary(node);
-    expect(typeof summary).toBe('string');
-    expect(summary.length).toBeGreaterThan(0);
-  });
-
-  it('handles non-English text', async () => {
-    const extractor = new SummaryExtractor();
-    const node = new TextNode({ text: '这是一个测试文档。' });
-    const summary = await extractor.generateNodeSummary(node);
-    expect(typeof summary).toBe('string');
-    expect(summary.length).toBeGreaterThan(0);
-  });
-
-  it('handles duplicate/repeated text', async () => {
-    const extractor = new SummaryExtractor();
-    const node = new TextNode({ text: 'repeat repeat repeat' });
-    const summary = await extractor.generateNodeSummary(node);
-    expect(typeof summary).toBe('string');
-    expect(summary.length).toBeGreaterThan(0);
-  });
-
-  it('handles only punctuation', async () => {
-    const extractor = new SummaryExtractor();
-    const node = new TextNode({ text: '!!!???...' });
-    const summary = await extractor.generateNodeSummary(node);
-    expect(typeof summary).toBe('string');
-    expect(summary.length).toBeGreaterThan(0);
-  });
-});

package/src/document/extractors/summary.ts

@@ -1,122 +0,0 @@
-import { Agent } from '@mastra/core/agent';
-import type { MastraLanguageModel } from '@mastra/core/agent';
-import { PromptTemplate, defaultSummaryPrompt } from '../prompts';
-import type { SummaryPrompt } from '../prompts';
-import type { BaseNode } from '../schema';
-import { TextNode } from '../schema';
-import { BaseExtractor } from './base';
-import { baseLLM, STRIP_REGEX } from './types';
-import type { SummaryExtractArgs } from './types';
-
-type ExtractSummary = {
-  sectionSummary?: string;
-  prevSectionSummary?: string;
-  nextSectionSummary?: string;
-};
-
-/**
- * Summarize an array of nodes using a custom LLM.
- *
- * @param nodes Array of node-like objects
- * @param options Summary extraction options
- * @returns Array of summary results
- */
-export class SummaryExtractor extends BaseExtractor {
-  private llm: MastraLanguageModel;
-  summaries: string[];
-  promptTemplate: SummaryPrompt;
-  private selfSummary: boolean;
-  private prevSummary: boolean;
-  private nextSummary: boolean;
-  constructor(options?: SummaryExtractArgs) {
-    const summaries = options?.summaries ?? ['self'];
-
-    if (summaries && !summaries.some(s => ['self', 'prev', 'next'].includes(s)))
-      throw new Error("Summaries must be one of 'self', 'prev', 'next'");
-
-    super();
-
-    this.llm = options?.llm ?? baseLLM;
-    this.summaries = summaries;
-    this.promptTemplate = options?.promptTemplate
-      ? new PromptTemplate({
-          templateVars: ['context'],
-          template: options.promptTemplate,
-        })
-      : defaultSummaryPrompt;
-
-    this.selfSummary = summaries?.includes('self') ?? false;
-    this.prevSummary = summaries?.includes('prev') ?? false;
-    this.nextSummary = summaries?.includes('next') ?? false;
-  }
-
-  /**
-   * Extract summary from a node.
-   * @param {BaseNode} node Node to extract summary from.
-   * @returns {Promise<string>} Summary extracted from the node.
-   */
-  async generateNodeSummary(node: BaseNode): Promise<string> {
-    const text = node.getContent();
-    if (!text || text.trim() === '') {
-      return '';
-    }
-    if (this.isTextNodeOnly && !(node instanceof TextNode)) {
-      return '';
-    }
-    const context = node.getContent();
-
-    const prompt = this.promptTemplate.format({
-      context,
-    });
-
-    const miniAgent = new Agent({
-      model: this.llm,
-      name: 'summary-extractor',
-      instructions:
-        'You are a summary extractor. You are given a node and you need to extract the summary from the node.',
-    });
-
-    let summary = '';
-    if (this.llm.specificationVersion === 'v2') {
-      const result = await miniAgent.generateVNext([{ role: 'user', content: prompt }], { format: 'mastra' });
-      summary = result.text;
-    } else {
-      const result = await miniAgent.generate([{ role: 'user', content: prompt }]);
-      summary = result.text;
-    }
-
-    if (!summary) {
-      console.warn('Summary extraction LLM output returned empty');
-      return '';
-    }
-
-    return summary.replace(STRIP_REGEX, '');
-  }
-
-  /**
-   * Extract summaries from a list of nodes.
-   * @param {BaseNode[]} nodes Nodes to extract summaries from.
-   * @returns {Promise<ExtractSummary[]>} Summaries extracted from the nodes.
-   */
-  async extract(nodes: BaseNode[]): Promise<ExtractSummary[]> {
-    if (!nodes.every(n => n instanceof TextNode)) throw new Error('Only `TextNode` is allowed for `Summary` extractor');
-
-    const nodeSummaries = await Promise.all(nodes.map(node => this.generateNodeSummary(node)));
-
-    const metadataList: ExtractSummary[] = nodes.map(() => ({}));
-
-    for (let i = 0; i < nodes.length; i++) {
-      if (i > 0 && this.prevSummary && nodeSummaries[i - 1]) {
-        metadataList[i]!['prevSectionSummary'] = nodeSummaries[i - 1];
-      }
-      if (i < nodes.length - 1 && this.nextSummary && nodeSummaries[i + 1]) {
-        metadataList[i]!['nextSectionSummary'] = nodeSummaries[i + 1];
-      }
-      if (this.selfSummary && nodeSummaries[i]) {
-        metadataList[i]!['sectionSummary'] = nodeSummaries[i];
-      }
-    }
-
-    return metadataList;
-  }
-}

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

@@ -1,121 +0,0 @@
-import { createOpenAI } from '@ai-sdk/openai';
-import { describe, it, expect, vi } from 'vitest';
-import { TextNode } from '../schema';
-import { TitleExtractor } from './title';
-
-const openai = createOpenAI({
-  apiKey: process.env.OPENAI_API_KEY,
-});
-
-const model = openai('gpt-4o');
-
-vi.setConfig({ testTimeout: 50_000, hookTimeout: 50_000 });
-
-describe('TitleExtractor', () => {
-  it('can use a custom model from the test suite', async () => {
-    const extractor = new TitleExtractor({ llm: model });
-    const node = new TextNode({ text: 'A title test using a custom model.' });
-    const titles = await extractor.extract([node]);
-    expect(Array.isArray(titles)).toBe(true);
-    expect(titles[0]).toHaveProperty('documentTitle');
-    expect(typeof titles[0].documentTitle).toBe('string');
-    expect(titles[0].documentTitle.length).toBeGreaterThan(0);
-  });
-
-  it('extracts title', async () => {
-    const extractor = new TitleExtractor({ llm: model });
-    const node = new TextNode({ text: 'This is a test document.' });
-    const titles = await extractor.extract([node]);
-    expect(Array.isArray(titles)).toBe(true);
-    expect(titles[0]).toHaveProperty('documentTitle');
-    expect(typeof titles[0].documentTitle).toBe('string');
-    expect(titles[0].documentTitle.length).toBeGreaterThan(0);
-  });
-
-  it('handles empty input gracefully', async () => {
-    const extractor = new TitleExtractor({ llm: model });
-    const node = new TextNode({ text: '' });
-    const titles = await extractor.extract([node]);
-    expect(titles[0].documentTitle).toBe('');
-  });
-
-  it('supports prompt customization', async () => {
-    const extractor = new TitleExtractor({ llm: model, nodeTemplate: 'Title for: {context}' });
-    const node = new TextNode({ text: 'Test document for prompt customization.' });
-    const titles = await extractor.extract([node]);
-    expect(titles[0]).toHaveProperty('documentTitle');
-    expect(typeof titles[0].documentTitle).toBe('string');
-    expect(titles[0].documentTitle.length).toBeGreaterThan(0);
-  });
-
-  it('handles very long input', async () => {
-    const extractor = new TitleExtractor({ llm: model });
-    const longText = 'A'.repeat(1000);
-    const node = new TextNode({ text: longText });
-    const titles = await extractor.extract([node]);
-    expect(titles[0]).toHaveProperty('documentTitle');
-    expect(typeof titles[0].documentTitle).toBe('string');
-    expect(titles[0].documentTitle.length).toBeGreaterThan(0);
-  });
-
-  it('handles whitespace only input', async () => {
-    const extractor = new TitleExtractor({ llm: model });
-    const node = new TextNode({ text: '    ' });
-    const titles = await extractor.extract([node]);
-    expect(titles[0].documentTitle).toBe('');
-  });
-
-  it('handles special characters and emojis', async () => {
-    const extractor = new TitleExtractor({ llm: model });
-    const node = new TextNode({ text: '🚀✨🔥' });
-    const titles = await extractor.extract([node]);
-    expect(titles[0]).toHaveProperty('documentTitle');
-    expect(typeof titles[0].documentTitle).toBe('string');
-    expect(titles[0].documentTitle.length).toBeGreaterThan(0);
-  });
-
-  it('handles numbers only', async () => {
-    const extractor = new TitleExtractor({ llm: model });
-    const node = new TextNode({ text: '1234567890' });
-    const titles = await extractor.extract([node]);
-    expect(titles[0]).toHaveProperty('documentTitle');
-    expect(typeof titles[0].documentTitle).toBe('string');
-    expect(titles[0].documentTitle.length).toBeGreaterThan(0);
-  });
-
-  it('handles HTML tags', async () => {
-    const extractor = new TitleExtractor({ llm: model });
-    const node = new TextNode({ text: '<h1>Test</h1>' });
-    const titles = await extractor.extract([node]);
-    expect(titles[0]).toHaveProperty('documentTitle');
-    expect(typeof titles[0].documentTitle).toBe('string');
-    expect(titles[0].documentTitle.length).toBeGreaterThan(0);
-  });
-
-  it('handles non-English text', async () => {
-    const extractor = new TitleExtractor({ llm: model });
-    const node = new TextNode({ text: '这是一个测试文档。' });
-    const titles = await extractor.extract([node]);
-    expect(titles[0]).toHaveProperty('documentTitle');
-    expect(typeof titles[0].documentTitle).toBe('string');
-    expect(titles[0].documentTitle.length).toBeGreaterThan(0);
-  });
-
-  it('handles duplicate/repeated text', async () => {
-    const extractor = new TitleExtractor({ llm: model });
-    const node = new TextNode({ text: 'repeat repeat repeat' });
-    const titles = await extractor.extract([node]);
-    expect(titles[0]).toHaveProperty('documentTitle');
-    expect(typeof titles[0].documentTitle).toBe('string');
-    expect(titles[0].documentTitle.length).toBeGreaterThan(0);
-  });
-
-  it('handles only punctuation', async () => {
-    const extractor = new TitleExtractor({ llm: model });
-    const node = new TextNode({ text: '!!!???...' });
-    const titles = await extractor.extract([node]);
-    expect(titles[0]).toHaveProperty('documentTitle');
-    expect(typeof titles[0].documentTitle).toBe('string');
-    expect(titles[0].documentTitle.length).toBeGreaterThan(0);
-  });
-});