@mastra/rag 1.0.8 → 1.0.9-alpha.0

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mastra/rag",
-  "version": "1.0.8",
+  "version": "1.0.9-alpha.0",
   "description": "",
   "type": "module",
   "main": "dist/index.js",
@@ -46,8 +46,8 @@
     "typescript": "^5.8.3",
     "vitest": "^3.2.4",
     "@internal/lint": "0.0.29",
-    "@mastra/core": "0.13.2",
-    "@internal/types-builder": "0.0.4"
+    "@internal/types-builder": "0.0.4",
+    "@mastra/core": "0.14.0-alpha.4"
   },
   "keywords": [
     "rag",

package/src/document/document.test.ts

@@ -20,7 +20,7 @@ const openai = createOpenAI({
   apiKey: process.env.OPENAI_API_KEY,
 });
 
-vi.setConfig({ testTimeout: 10_000, hookTimeout: 10_000 });
+vi.setConfig({ testTimeout: 100_000, hookTimeout: 100_000 });
 
 describe('MDocument', () => {
   describe('basics', () => {

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

@@ -9,7 +9,7 @@ const openai = createOpenAI({
 
 const model = openai('gpt-4o');
 
-vi.setConfig({ testTimeout: 10_000, hookTimeout: 10_000 });
+vi.setConfig({ testTimeout: 50_000, hookTimeout: 50_000 });
 
 describe('KeywordExtractor', () => {
   it('can use a custom model for keywords extraction', async () => {
@@ -49,7 +49,7 @@ describe('KeywordExtractor', () => {
   it(
     'handles very long input',
     {
-      timeout: 20_000,
+      timeout: 60_000,
     },
     async () => {
       const extractor = new KeywordExtractor();

package/src/document/extractors/keywords.ts

@@ -1,3 +1,4 @@
+import { Agent } from '@mastra/core/agent';
 import type { MastraLanguageModel } from '@mastra/core/agent';
 import { defaultKeywordExtractPrompt, PromptTemplate } from '../prompts';
 import type { KeywordExtractPrompt } from '../prompts';
@@ -64,33 +65,47 @@ export class KeywordExtractor extends BaseExtractor {
 
     let keywords = '';
     try {
-      const completion = await this.llm.doGenerate({
-        inputFormat: 'messages',
-        mode: { type: 'regular' },
-        prompt: [
+      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: [
-              {
-                type: 'text',
-                text: this.promptTemplate.format({
-                  context: node.getContent(),
-                  maxKeywords: this.keywords.toString(),
-                }),
-              },
-            ],
+            content: this.promptTemplate.format({ context: node.getContent(), maxKeywords: this.keywords.toString() }),
           },
-        ],
-      });
-      if (typeof completion.text === 'string') {
-        keywords = completion.text.trim();
-      } else {
-        console.warn('Keyword extraction LLM output was not a string:', completion.text);
+        ]);
+        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: '' };
     }
-    return { excerptKeywords: keywords };
   }
 
   /**

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

@@ -9,7 +9,7 @@ const openai = createOpenAI({
 
 const model = openai('gpt-4o');
 
-vi.setConfig({ testTimeout: 10_000, hookTimeout: 10_000 });
+vi.setConfig({ testTimeout: 100_000, hookTimeout: 100_000 });
 
 describe('QuestionsAnsweredExtractor', () => {
   it('can use a custom model for questions extraction', async () => {

package/src/document/extractors/questions.ts

@@ -1,3 +1,4 @@
+import { Agent } from '@mastra/core/agent';
 import type { MastraLanguageModel } from '@mastra/core/agent';
 import { PromptTemplate, defaultQuestionExtractPrompt } from '../prompts';
 import type { QuestionExtractPrompt } from '../prompts';
@@ -69,27 +70,29 @@ export class QuestionsAnsweredExtractor extends BaseExtractor {
       numQuestions: this.questions.toString(),
     });
 
-    const questions = await this.llm.doGenerate({
-      inputFormat: 'messages',
-      mode: { type: 'regular' },
-      prompt: [
-        {
-          role: 'user',
-          content: [{ type: 'text', text: prompt }],
-        },
-      ],
+    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 result = '';
-    try {
-      if (typeof questions.text === 'string') {
-        result = questions.text.replace(STRIP_REGEX, '').trim();
-      } else {
-        console.warn('Question extraction LLM output was not a string:', questions.text);
-      }
-    } catch (err) {
-      console.warn('Question extraction failed:', err);
+    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,
     };

package/src/document/extractors/summary.ts

@@ -1,3 +1,4 @@
+import { Agent } from '@mastra/core/agent';
 import type { MastraLanguageModel } from '@mastra/core/agent';
 import { PromptTemplate, defaultSummaryPrompt } from '../prompts';
 import type { SummaryPrompt } from '../prompts';
@@ -68,22 +69,25 @@ export class SummaryExtractor extends BaseExtractor {
       context,
     });
 
-    const result = await this.llm.doGenerate({
-      inputFormat: 'messages',
-      mode: { type: 'regular' },
-      prompt: [
-        {
-          role: 'user',
-          content: [{ type: 'text', text: prompt }],
-        },
-      ],
+    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 (typeof result.text === 'string') {
-      summary = result.text.trim();
+    if (this.llm.specificationVersion === 'v2') {
+      const result = await miniAgent.generateVNext([{ role: 'user', content: prompt }], { format: 'mastra' });
+      summary = result.text;
     } else {
-      console.warn('Summary extraction LLM output was not a string:', result.text);
+      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, '');

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

@@ -9,7 +9,7 @@ const openai = createOpenAI({
 
 const model = openai('gpt-4o');
 
-vi.setConfig({ testTimeout: 10_000, hookTimeout: 10_000 });
+vi.setConfig({ testTimeout: 50_000, hookTimeout: 50_000 });
 
 describe('TitleExtractor', () => {
   it('can use a custom model from the test suite', async () => {

package/src/document/extractors/title.ts

@@ -1,3 +1,4 @@
+import { Agent } from '@mastra/core/agent';
 import type { MastraLanguageModel } from '@mastra/core/agent';
 import { defaultTitleCombinePromptTemplate, defaultTitleExtractorPromptTemplate, PromptTemplate } from '../prompts';
 import type { TitleCombinePrompt, TitleExtractorPrompt } from '../prompts';
@@ -110,30 +111,38 @@ export class TitleExtractor extends BaseExtractor {
     for (const [key, nodes] of Object.entries(nodesByDocument)) {
       const titleCandidates = await this.getTitlesCandidates(nodes);
       const combinedTitles = titleCandidates.join(', ');
-      const completion = await this.llm.doGenerate({
-        inputFormat: 'messages',
-        mode: { type: 'regular' },
-        prompt: [
-          {
-            role: 'user',
-            content: [
-              {
-                type: 'text',
-                text: this.combineTemplate.format({
-                  context: combinedTitles,
-                }),
-              },
-            ],
-          },
-        ],
-      });
 
       let title = '';
-      if (typeof completion.text === 'string') {
-        title = completion.text.trim();
+
+      if (this.llm.specificationVersion === 'v2') {
+        const miniAgent = new Agent({
+          model: this.llm,
+          name: 'title-extractor',
+          instructions:
+            'You are a title extractor. You are given a list of nodes and you need to extract the title from the nodes.',
+        });
+        const result = await miniAgent.generateVNext(
+          [{ role: 'user', content: this.combineTemplate.format({ context: combinedTitles }) }],
+          { format: 'mastra' },
+        );
+        title = result.text;
       } else {
-        console.warn('Title extraction LLM output was not a string:', completion.text);
+        const miniAgent = new Agent({
+          model: this.llm,
+          name: 'title-extractor',
+          instructions:
+            'You are a title extractor. You are given a list of nodes and you need to extract the title from the nodes.',
+        });
+        const result = await miniAgent.generate([
+          { role: 'user', content: this.combineTemplate.format({ context: combinedTitles }) },
+        ]);
+        title = result.text;
       }
+
+      if (!title) {
+        console.warn('Title extraction LLM output returned empty');
+      }
+
       titlesByDocument[key] = title;
     }
 
@@ -141,31 +150,34 @@ export class TitleExtractor extends BaseExtractor {
   }
 
   private async getTitlesCandidates(nodes: BaseNode[]): Promise<string[]> {
+    const miniAgent = new Agent({
+      model: this.llm,
+      name: 'titles-candidates-extractor',
+      instructions:
+        'You are a titles candidates extractor. You are given a list of nodes and you need to extract the titles candidates from the nodes.',
+    });
+
     const titleJobs = nodes.map(async node => {
-      const completion = await this.llm.doGenerate({
-        inputFormat: 'messages',
-        mode: { type: 'regular' },
-        prompt: [
-          {
-            role: 'user',
-            content: [
-              {
-                type: 'text',
-                text: this.nodeTemplate.format({
-                  context: node.getContent(),
-                }),
-              },
-            ],
-          },
-        ],
-      });
-
-      if (typeof completion.text === 'string') {
-        return completion.text.trim();
+      let completion: string;
+      if (this.llm.specificationVersion === 'v2') {
+        const result = await miniAgent.generateVNext(
+          [{ role: 'user', content: this.nodeTemplate.format({ context: node.getContent() }) }],
+          { format: 'mastra' },
+        );
+        completion = result.text;
       } else {
-        console.warn('Title candidate extraction LLM output was not a string:', completion.text);
+        const result = await miniAgent.generate([
+          { role: 'user', content: this.nodeTemplate.format({ context: node.getContent() }) },
+        ]);
+        completion = result.text;
+      }
+
+      if (!completion) {
+        console.warn('Title candidate extraction LLM output returned empty');
         return '';
       }
+
+      return completion.trim();
     });
 
     return await Promise.all(titleJobs);