@mastra/rag 0.1.19 → 0.1.20-alpha.1

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@mastra/rag",
-  "version": "0.1.19",
+  "version": "0.1.20-alpha.1",
   "description": "",
   "type": "module",
   "main": "dist/index.js",
@@ -23,11 +23,10 @@
   "dependencies": {
     "@paralleldrive/cuid2": "^2.2.2",
     "js-tiktoken": "^1.0.19",
-    "llamaindex": "^0.9.17",
     "node-html-better-parser": "^1.4.7",
     "pathe": "^2.0.3",
     "zod": "^3.24.2",
-    "@mastra/core": "^0.9.0"
+    "@mastra/core": "^0.9.1-alpha.1"
   },
   "peerDependencies": {
     "ai": "^4.0.0"

package/src/document/document.ts

@@ -1,6 +1,6 @@
-import { Document as Chunk, IngestionPipeline, NodeRelationship, ObjectType } from 'llamaindex';
-
 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';
@@ -54,13 +54,10 @@ export class MDocument {
       );
     }
 
-    const pipeline = new IngestionPipeline({
-      transformations,
-    });
-
-    const nodes = await pipeline.run({
-      documents: this.chunks,
-    });
+    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({

package/src/document/extractors/base.ts

@@ -0,0 +1,30 @@
+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

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

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

@@ -1,6 +1,6 @@
 import { createOpenAI } from '@ai-sdk/openai';
-import { TextNode } from 'llamaindex';
 import { describe, it, expect, vi } from 'vitest';
+import { TextNode } from '../schema';
 import { KeywordExtractor } from './keywords';
 
 const openai = createOpenAI({

package/src/document/extractors/keywords.ts

@@ -1,6 +1,9 @@
 import type { MastraLanguageModel } from '@mastra/core/agent';
-import { PromptTemplate, defaultKeywordExtractPrompt, MetadataMode, TextNode, BaseExtractor } from 'llamaindex';
-import type { KeywordExtractPrompt, BaseNode } from 'llamaindex';
+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';
 
@@ -15,23 +18,8 @@ type ExtractKeyword = {
  * Extract keywords from a list of nodes.
  */
 export class KeywordExtractor extends BaseExtractor {
-  /**
-   * MastraLanguageModel instance.
-   * @type {MastraLanguageModel}
-   */
   llm: MastraLanguageModel;
-
-  /**
-   * Number of keywords to extract.
-   * @type {number}
-   * @default 5
-   */
   keywords: number = 5;
-
-  /**
-   * The prompt template to use for the question extractor.
-   * @type {string}
-   */
   promptTemplate: KeywordExtractPrompt;
 
   /**
@@ -66,7 +54,7 @@ export class KeywordExtractor extends BaseExtractor {
    * Adds error handling for malformed/empty LLM output.
    */
   async extractKeywordsFromNodes(node: BaseNode): Promise<ExtractKeyword> {
-    const text = node.getContent(this.metadataMode);
+    const text = node.getContent();
     if (!text || text.trim() === '') {
       return { excerptKeywords: '' };
     }
@@ -86,7 +74,7 @@ export class KeywordExtractor extends BaseExtractor {
               {
                 type: 'text',
                 text: this.promptTemplate.format({
-                  context: node.getContent(MetadataMode.ALL),
+                  context: node.getContent(),
                   maxKeywords: this.keywords.toString(),
                 }),
               },

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

@@ -1,6 +1,6 @@
 import { createOpenAI } from '@ai-sdk/openai';
-import { TextNode } from 'llamaindex';
 import { describe, it, expect, vi } from 'vitest';
+import { TextNode } from '../schema';
 import { QuestionsAnsweredExtractor } from './questions';
 
 const openai = createOpenAI({

package/src/document/extractors/questions.ts

@@ -1,6 +1,9 @@
 import type { MastraLanguageModel } from '@mastra/core/agent';
-import { PromptTemplate, defaultQuestionExtractPrompt, TextNode, BaseExtractor } from 'llamaindex';
-import type { QuestionExtractPrompt, BaseNode } from 'llamaindex';
+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';
 
@@ -15,30 +18,9 @@ type ExtractQuestion = {
  * Extract questions from a list of nodes.
  */
 export class QuestionsAnsweredExtractor extends BaseExtractor {
-  /**
-   * MastraLanguageModel instance.
-   * @type {MastraLanguageModel}
-   */
   llm: MastraLanguageModel;
-
-  /**
-   * Number of questions to generate.
-   * @type {number}
-   * @default 5
-   */
   questions: number = 5;
-
-  /**
-   * The prompt template to use for the question extractor.
-   * @type {string}
-   */
   promptTemplate: QuestionExtractPrompt;
-
-  /**
-   * Wheter to use metadata for embeddings only
-   * @type {boolean}
-   * @default false
-   */
   embeddingOnly: boolean = false;
 
   /**
@@ -72,7 +54,7 @@ export class QuestionsAnsweredExtractor extends BaseExtractor {
    * @returns {Promise<Array<ExtractQuestion> | Array<{}>>} Questions extracted from the node.
    */
   async extractQuestionsFromNode(node: BaseNode): Promise<ExtractQuestion> {
-    const text = node.getContent(this.metadataMode);
+    const text = node.getContent();
     if (!text || text.trim() === '') {
       return { questionsThisExcerptCanAnswer: '' };
     }
@@ -80,7 +62,7 @@ export class QuestionsAnsweredExtractor extends BaseExtractor {
       return { questionsThisExcerptCanAnswer: '' };
     }
 
-    const contextStr = node.getContent(this.metadataMode);
+    const contextStr = node.getContent();
 
     const prompt = this.promptTemplate.format({
       context: contextStr,

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

@@ -1,6 +1,6 @@
 import { createOpenAI } from '@ai-sdk/openai';
-import { TextNode } from 'llamaindex';
 import { describe, it, expect, vi } from 'vitest';
+import { TextNode } from '../schema';
 import { SummaryExtractor } from './summary';
 
 const openai = createOpenAI({

package/src/document/extractors/summary.ts

@@ -1,6 +1,9 @@
 import type { MastraLanguageModel } from '@mastra/core/agent';
-import { PromptTemplate, defaultSummaryPrompt, TextNode, BaseExtractor } from 'llamaindex';
-import type { SummaryPrompt, BaseNode } from 'llamaindex';
+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';
 
@@ -18,27 +21,12 @@ type ExtractSummary = {
  * @returns Array of summary results
  */
 export class SummaryExtractor extends BaseExtractor {
-  /**
-   * MastraLanguageModel instance.
-   * @type {MastraLanguageModel}
-   */
   private llm: MastraLanguageModel;
-  /**
-   * List of summaries to extract: 'self', 'prev', 'next'
-   * @type {string[]}
-   */
   summaries: string[];
-
-  /**
-   * The prompt template to use for the summary extractor.
-   * @type {string}
-   */
   promptTemplate: SummaryPrompt;
-
   private selfSummary: boolean;
   private prevSummary: boolean;
   private nextSummary: boolean;
-
   constructor(options?: SummaryExtractArgs) {
     const summaries = options?.summaries ?? ['self'];
 
@@ -67,14 +55,14 @@ export class SummaryExtractor extends BaseExtractor {
    * @returns {Promise<string>} Summary extracted from the node.
    */
   async generateNodeSummary(node: BaseNode): Promise<string> {
-    const text = node.getContent(this.metadataMode);
+    const text = node.getContent();
     if (!text || text.trim() === '') {
       return '';
     }
     if (this.isTextNodeOnly && !(node instanceof TextNode)) {
       return '';
     }
-    const context = node.getContent(this.metadataMode);
+    const context = node.getContent();
 
     const prompt = this.promptTemplate.format({
       context,

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

@@ -1,6 +1,6 @@
 import { createOpenAI } from '@ai-sdk/openai';
-import { TextNode } from 'llamaindex';
 import { describe, it, expect, vi } from 'vitest';
+import { TextNode } from '../schema';
 import { TitleExtractor } from './title';
 
 const openai = createOpenAI({

package/src/document/extractors/title.ts

@@ -1,13 +1,9 @@
 import type { MastraLanguageModel } from '@mastra/core/agent';
-import {
-  PromptTemplate,
-  defaultTitleCombinePromptTemplate,
-  defaultTitleExtractorPromptTemplate,
-  MetadataMode,
-  TextNode,
-  BaseExtractor,
-} from 'llamaindex';
-import type { TitleCombinePrompt, TitleExtractorPrompt, BaseNode } from 'llamaindex';
+import { defaultTitleCombinePromptTemplate, defaultTitleExtractorPromptTemplate, PromptTemplate } from '../prompts';
+import type { TitleCombinePrompt, TitleExtractorPrompt } from '../prompts';
+import { TextNode } from '../schema';
+import type { BaseNode } from '../schema';
+import { BaseExtractor } from './base';
 import { baseLLM } from './types';
 import type { TitleExtractorsArgs } from './types';
 
@@ -19,45 +15,12 @@ type ExtractTitle = {
  * Extract title from a list of nodes.
  */
 export class TitleExtractor extends BaseExtractor {
-  /**
-   * MastraLanguageModel instance.
-   * @type {MastraLanguageModel}
-   */
   llm: MastraLanguageModel;
-
-  /**
-   * Can work for mixture of text and non-text nodes
-   * @type {boolean}
-   * @default false
-   */
   isTextNodeOnly: boolean = false;
-
-  /**
-   * Number of nodes to extrct titles from.
-   * @type {number}
-   * @default 5
-   */
   nodes: number = 5;
-
-  /**
-   * The prompt template to use for the title extractor.
-   * @type {string}
-   */
   nodeTemplate: TitleExtractorPrompt;
-
-  /**
-   * The prompt template to merge title with..
-   * @type {string}
-   */
   combineTemplate: TitleCombinePrompt;
 
-  /**
-   * Constructor for the TitleExtractor class.
-   * @param {MastraLanguageModel} llm MastraLanguageModel instance.
-   * @param {number} nodes Number of nodes to extract titles from.
-   * @param {TitleExtractorPrompt} nodeTemplate The prompt template to use for the title extractor.
-   * @param {string} combineTemplate The prompt template to merge title with..
-   */
   constructor(options?: TitleExtractorsArgs) {
     super();
 
@@ -92,7 +55,7 @@ export class TitleExtractor extends BaseExtractor {
     const nodeIndexes: number[] = [];
 
     nodes.forEach((node, idx) => {
-      const text = node.getContent(this.metadataMode);
+      const text = node.getContent();
       if (!text || text.trim() === '') {
         results[idx] = { documentTitle: '' };
       } else {
@@ -189,7 +152,7 @@ export class TitleExtractor extends BaseExtractor {
               {
                 type: 'text',
                 text: this.nodeTemplate.format({
-                  context: node.getContent(MetadataMode.ALL),
+                  context: node.getContent(),
                 }),
               },
             ],

package/src/document/extractors/types.ts

@@ -6,7 +6,7 @@ import type {
   SummaryPrompt,
   TitleExtractorPrompt,
   TitleCombinePrompt,
-} from 'llamaindex';
+} from '../prompts';
 
 export type KeywordExtractArgs = {
   llm?: MastraLanguageModel;

package/src/document/prompts/base.ts

@@ -0,0 +1,77 @@
+import { format } from './format';
+import type { BasePromptTemplateOptions, ChatMessage, PromptTemplateOptions } from './types';
+
+export abstract class BasePromptTemplate<const TemplatesVar extends readonly string[] = string[]> {
+  templateVars: Set<string> = new Set();
+  options: Partial<Record<TemplatesVar[number] | (string & {}), string>> = {};
+
+  protected constructor(options: BasePromptTemplateOptions<TemplatesVar>) {
+    const { templateVars } = options;
+    if (templateVars) {
+      this.templateVars = new Set(templateVars);
+    }
+    if (options.options) {
+      this.options = options.options;
+    }
+  }
+
+  abstract partialFormat(
+    options: Partial<Record<TemplatesVar[number] | (string & {}), string>>,
+  ): BasePromptTemplate<TemplatesVar>;
+
+  abstract format(options?: Partial<Record<TemplatesVar[number] | (string & {}), string>>): string;
+
+  abstract formatMessages(options?: Partial<Record<TemplatesVar[number] | (string & {}), string>>): ChatMessage[];
+
+  abstract get template(): string;
+}
+
+export class PromptTemplate<
+  const TemplatesVar extends readonly string[] = string[],
+> extends BasePromptTemplate<TemplatesVar> {
+  #template: string;
+
+  constructor(options: PromptTemplateOptions<TemplatesVar>) {
+    const { template, ...rest } = options;
+    super(rest);
+    this.#template = template;
+  }
+
+  partialFormat(options: Partial<Record<TemplatesVar[number] | (string & {}), string>>): PromptTemplate<TemplatesVar> {
+    const prompt = new PromptTemplate({
+      template: this.template,
+      templateVars: [...this.templateVars],
+      options: this.options,
+    });
+
+    prompt.options = {
+      ...prompt.options,
+      ...options,
+    };
+
+    return prompt;
+  }
+
+  format(options?: Partial<Record<TemplatesVar[number] | (string & {}), string>>): string {
+    const allOptions = {
+      ...this.options,
+      ...options,
+    } as Record<TemplatesVar[number], string>;
+
+    return format(this.template, allOptions);
+  }
+
+  formatMessages(options?: Partial<Record<TemplatesVar[number] | (string & {}), string>>): ChatMessage[] {
+    const prompt = this.format(options);
+    return [
+      {
+        role: 'user',
+        content: prompt,
+      },
+    ];
+  }
+
+  get template(): string {
+    return this.#template;
+  }
+}

package/src/document/prompts/format.ts

@@ -0,0 +1,9 @@
+/**
+ * Formats a string by replacing placeholders with values from the provided parameters.
+ * @param str The string to format.
+ * @param params A record of placeholder names to their corresponding values.
+ * @returns The formatted string.
+ */
+export function format(str: string, params: Record<string, string>) {
+  return str.replace(/{(\w+)}/g, (_, k) => params[k] ?? '');
+}

package/src/document/prompts/index.ts

@@ -0,0 +1,15 @@
+export { BasePromptTemplate, PromptTemplate } from './base';
+export {
+  defaultKeywordExtractPrompt,
+  defaultQuestionExtractPrompt,
+  defaultSummaryPrompt,
+  defaultTitleCombinePromptTemplate,
+  defaultTitleExtractorPromptTemplate,
+} from './prompt';
+export type {
+  KeywordExtractPrompt,
+  QuestionExtractPrompt,
+  SummaryPrompt,
+  TitleCombinePrompt,
+  TitleExtractorPrompt,
+} from './prompt';

package/src/document/prompts/prompt.ts

@@ -0,0 +1,60 @@
+import { PromptTemplate } from './base';
+
+export type SummaryPrompt = PromptTemplate<['context']>;
+export type KeywordExtractPrompt = PromptTemplate<['context', 'maxKeywords']>;
+export type QuestionExtractPrompt = PromptTemplate<['context', 'numQuestions']>;
+export type TitleExtractorPrompt = PromptTemplate<['context']>;
+export type TitleCombinePrompt = PromptTemplate<['context']>;
+
+export const defaultSummaryPrompt: SummaryPrompt = new PromptTemplate({
+  templateVars: ['context'],
+  template: `Write a summary of the following. Try to use only the information provided. Try to include as many key details as possible.
+
+
+{context}
+
+
+SUMMARY:"""
+`,
+});
+
+export const defaultKeywordExtractPrompt: KeywordExtractPrompt = new PromptTemplate({
+  templateVars: ['maxKeywords', 'context'],
+  template: `
+Some text is provided below. Given the text, extract up to {maxKeywords} keywords from the text. Avoid stopwords.
+---------------------
+{context}
+---------------------
+Provide keywords in the following comma-separated format: 'KEYWORDS: <keywords>'
+`,
+}).partialFormat({
+  maxKeywords: '10',
+});
+
+export const defaultQuestionExtractPrompt = new PromptTemplate({
+  templateVars: ['numQuestions', 'context'],
+  template: `(
+  "Given the contextual informations below, generate {numQuestions} questions this context can provides specific answers to which are unlikely to be found else where. Higher-level summaries of surrounding context may be provided as well. "
+  "Try using these summaries to generate better questions that this context can answer."
+  "---------------------"
+  "{context}"
+  "---------------------"
+  "Provide questions in the following format: 'QUESTIONS: <questions>'"
+)`,
+}).partialFormat({
+  numQuestions: '5',
+});
+
+export const defaultTitleExtractorPromptTemplate = new PromptTemplate({
+  templateVars: ['context'],
+  template: `{context}
+Give a title that summarizes all of the unique entities, titles or themes found in the context. 
+Title: `,
+});
+
+export const defaultTitleCombinePromptTemplate = new PromptTemplate({
+  templateVars: ['context'],
+  template: `{context} 
+Based on the above candidate titles and contents, what is the comprehensive title for this document? 
+Title: `,
+});

package/src/document/prompts/types.ts

@@ -0,0 +1,29 @@
+export type MessageType = 'user' | 'assistant' | 'system' | 'memory' | 'developer';
+
+export type MessageContentTextDetail = {
+  type: 'text';
+  text: string;
+};
+
+/**
+ * Extended type for the content of a message that allows for multi-modal messages.
+ */
+export type MessageContent = string | MessageContentTextDetail[];
+
+export type ChatMessage<AdditionalMessageOptions extends object = object> = {
+  content: MessageContent;
+  role: MessageType;
+  options?: undefined | AdditionalMessageOptions;
+};
+
+export type BasePromptTemplateOptions<TemplatesVar extends readonly string[]> = {
+  templateVars?:
+    | TemplatesVar
+    // loose type for better type inference
+    | readonly string[];
+  options?: Partial<Record<TemplatesVar[number] | (string & {}), string>>;
+};
+
+export type PromptTemplateOptions<TemplatesVar extends readonly string[]> = BasePromptTemplateOptions<TemplatesVar> & {
+  template: string;
+};

package/src/document/schema/index.ts

@@ -0,0 +1,3 @@
+export { BaseNode, Document, TextNode } from './node';
+export { NodeRelationship, ObjectType } from './types';
+export type { Metadata, RelatedNodeInfo, RelatedNodeType, BaseNodeParams, TextNodeParams } from './types';