@mastra/rag 1.2.3-alpha.0 → 1.2.3-alpha.2

package/src/document/extractors/title.ts

@@ -1,185 +0,0 @@
-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';
-import { TextNode } from '../schema';
-import type { BaseNode } from '../schema';
-import { BaseExtractor } from './base';
-import { baseLLM } from './types';
-import type { TitleExtractorsArgs } from './types';
-
-type ExtractTitle = {
-  documentTitle: string;
-};
-
-/**
- * Extract title from a list of nodes.
- */
-export class TitleExtractor extends BaseExtractor {
-  llm: MastraLanguageModel;
-  isTextNodeOnly: boolean = false;
-  nodes: number = 5;
-  nodeTemplate: TitleExtractorPrompt;
-  combineTemplate: TitleCombinePrompt;
-
-  constructor(options?: TitleExtractorsArgs) {
-    super();
-
-    this.llm = options?.llm ?? baseLLM;
-    this.nodes = options?.nodes ?? 5;
-
-    this.nodeTemplate = options?.nodeTemplate
-      ? new PromptTemplate({
-          templateVars: ['context'],
-          template: options.nodeTemplate,
-        })
-      : defaultTitleExtractorPromptTemplate;
-
-    this.combineTemplate = options?.combineTemplate
-      ? new PromptTemplate({
-          templateVars: ['context'],
-          template: options.combineTemplate,
-        })
-      : defaultTitleCombinePromptTemplate;
-  }
-
-  /**
-   * Extract titles from a list of nodes.
-   * @param {BaseNode[]} nodes Nodes to extract titles from.
-   * @returns {Promise<BaseNode<ExtractTitle>[]>} Titles extracted from the nodes.
-   */
-  async extract(nodes: BaseNode[]): Promise<Array<ExtractTitle>> {
-    // Prepare output array in original node order
-    const results: ExtractTitle[] = new Array(nodes.length);
-    // Keep track of nodes with content to extract
-    const nodesToExtractTitle: BaseNode[] = [];
-    const nodeIndexes: number[] = [];
-
-    nodes.forEach((node, idx) => {
-      const text = node.getContent();
-      if (!text || text.trim() === '') {
-        results[idx] = { documentTitle: '' };
-      } else {
-        nodesToExtractTitle.push(node);
-        nodeIndexes.push(idx);
-      }
-    });
-
-    if (nodesToExtractTitle.length) {
-      const filteredNodes = this.filterNodes(nodesToExtractTitle);
-      if (filteredNodes.length) {
-        const nodesByDocument = this.separateNodesByDocument(filteredNodes);
-        const titlesByDocument = await this.extractTitles(nodesByDocument);
-        filteredNodes.forEach((node, i) => {
-          const nodeIndex = nodeIndexes[i];
-          const groupKey = node.sourceNode?.nodeId ?? node.id_;
-          if (typeof nodeIndex === 'number') {
-            results[nodeIndex] = {
-              documentTitle: titlesByDocument[groupKey] ?? '',
-            };
-          }
-        });
-      }
-    }
-    return results;
-  }
-
-  private filterNodes(nodes: BaseNode[]): BaseNode[] {
-    return nodes.filter(node => {
-      if (this.isTextNodeOnly && !(node instanceof TextNode)) {
-        return false;
-      }
-      return true;
-    });
-  }
-
-  private separateNodesByDocument(nodes: BaseNode[]): Record<string, BaseNode[]> {
-    const nodesByDocument: Record<string, BaseNode[]> = {};
-
-    for (const node of nodes) {
-      const groupKey = node.sourceNode?.nodeId ?? node.id_;
-      nodesByDocument[groupKey] = nodesByDocument[groupKey] || [];
-      nodesByDocument[groupKey].push(node);
-    }
-
-    return nodesByDocument;
-  }
-
-  private async extractTitles(nodesByDocument: Record<string, BaseNode[]>): Promise<Record<string, string>> {
-    const titlesByDocument: Record<string, string> = {};
-
-    for (const [key, nodes] of Object.entries(nodesByDocument)) {
-      const titleCandidates = await this.getTitlesCandidates(nodes);
-      const combinedTitles = titleCandidates.join(', ');
-
-      let title = '';
-
-      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 {
-        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;
-    }
-
-    return titlesByDocument;
-  }
-
-  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 => {
-      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 {
-        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);
-  }
-}

package/src/document/extractors/types.ts

@@ -1,40 +0,0 @@
-import { createOpenAI } from '@ai-sdk/openai';
-import type { MastraLanguageModel } from '@mastra/core/agent';
-import type {
-  KeywordExtractPrompt,
-  QuestionExtractPrompt,
-  SummaryPrompt,
-  TitleExtractorPrompt,
-  TitleCombinePrompt,
-} from '../prompts';
-
-export type KeywordExtractArgs = {
-  llm?: MastraLanguageModel;
-  keywords?: number;
-  promptTemplate?: KeywordExtractPrompt['template'];
-};
-
-export type QuestionAnswerExtractArgs = {
-  llm?: MastraLanguageModel;
-  questions?: number;
-  promptTemplate?: QuestionExtractPrompt['template'];
-  embeddingOnly?: boolean;
-};
-
-export type SummaryExtractArgs = {
-  llm?: MastraLanguageModel;
-  summaries?: string[];
-  promptTemplate?: SummaryPrompt['template'];
-};
-
-export type TitleExtractorsArgs = {
-  llm?: MastraLanguageModel;
-  nodes?: number;
-  nodeTemplate?: TitleExtractorPrompt['template'];
-  combineTemplate?: TitleCombinePrompt['template'];
-};
-
-export const STRIP_REGEX = /(\r\n|\n|\r)/gm;
-
-const openai = createOpenAI({ apiKey: process.env.OPENAI_API_KEY });
-export const baseLLM: MastraLanguageModel = openai('gpt-4o');

package/src/document/index.ts

@@ -1,2 +0,0 @@
-export * from './document';
-export * from './types';

package/src/document/prompts/base.ts

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

@@ -1,9 +0,0 @@
-/**
- * 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

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

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

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

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

package/src/document/schema/node.ts

@@ -1,187 +0,0 @@
-import { createHash, randomUUID } from 'crypto';
-import { NodeRelationship, ObjectType } from './types';
-import type { Metadata, RelatedNodeInfo, RelatedNodeType, BaseNodeParams, TextNodeParams } from './types';
-
-/**
- * Generic abstract class for retrievable nodes
- */
-export abstract class BaseNode<T extends Metadata = Metadata> {
-  id_: string;
-  metadata: T;
-  relationships: Partial<Record<NodeRelationship, RelatedNodeType<T>>>;
-
-  @lazyInitHash
-  accessor hash: string = '';
-
-  protected constructor(init?: BaseNodeParams<T>) {
-    const { id_, metadata, relationships } = init || {};
-    this.id_ = id_ ?? randomUUID();
-    this.metadata = metadata ?? ({} as T);
-    this.relationships = relationships ?? {};
-  }
-
-  abstract get type(): ObjectType;
-
-  abstract getContent(): string;
-
-  abstract getMetadataStr(): string;
-
-  get sourceNode(): RelatedNodeInfo<T> | undefined {
-    const relationship = this.relationships[NodeRelationship.SOURCE];
-
-    if (Array.isArray(relationship)) {
-      throw new Error('Source object must be a single RelatedNodeInfo object');
-    }
-
-    return relationship;
-  }
-
-  get prevNode(): RelatedNodeInfo<T> | undefined {
-    const relationship = this.relationships[NodeRelationship.PREVIOUS];
-
-    if (Array.isArray(relationship)) {
-      throw new Error('Previous object must be a single RelatedNodeInfo object');
-    }
-
-    return relationship;
-  }
-
-  get nextNode(): RelatedNodeInfo<T> | undefined {
-    const relationship = this.relationships[NodeRelationship.NEXT];
-
-    if (Array.isArray(relationship)) {
-      throw new Error('Next object must be a single RelatedNodeInfo object');
-    }
-
-    return relationship;
-  }
-
-  get parentNode(): RelatedNodeInfo<T> | undefined {
-    const relationship = this.relationships[NodeRelationship.PARENT];
-
-    if (Array.isArray(relationship)) {
-      throw new Error('Parent object must be a single RelatedNodeInfo object');
-    }
-
-    return relationship;
-  }
-
-  get childNodes(): RelatedNodeInfo<T>[] | undefined {
-    const relationship = this.relationships[NodeRelationship.CHILD];
-
-    if (!Array.isArray(relationship)) {
-      throw new Error('Child object must be a an array of RelatedNodeInfo objects');
-    }
-
-    return relationship;
-  }
-
-  abstract generateHash(): string;
-}
-
-/**
- * TextNode is the default node type for text.
- */
-export class TextNode<T extends Metadata = Metadata> extends BaseNode<T> {
-  text: string;
-
-  startCharIdx?: number;
-  endCharIdx?: number;
-  metadataSeparator: string;
-
-  constructor(init: TextNodeParams<T> = {}) {
-    super(init);
-    const { text, startCharIdx, endCharIdx, metadataSeparator } = init;
-    this.text = text ?? '';
-    if (startCharIdx) {
-      this.startCharIdx = startCharIdx;
-    }
-    if (endCharIdx) {
-      this.endCharIdx = endCharIdx;
-    }
-    this.metadataSeparator = metadataSeparator ?? '\n';
-  }
-
-  /**
-   * Generate a hash of the text node.
-   * The ID is not part of the hash as it can change independent of content.
-   * @returns
-   */
-  generateHash() {
-    const hashFunction = createSHA256();
-    hashFunction.update(`type=${this.type}`);
-    hashFunction.update(`startCharIdx=${this.startCharIdx} endCharIdx=${this.endCharIdx}`);
-    hashFunction.update(this.getContent());
-    return hashFunction.digest();
-  }
-
-  get type() {
-    return ObjectType.TEXT;
-  }
-
-  getContent(): string {
-    const metadataStr = this.getMetadataStr().trim();
-    return `${metadataStr}\n\n${this.text}`.trim();
-  }
-
-  getMetadataStr(): string {
-    const usableMetadataKeys = new Set(Object.keys(this.metadata).sort());
-
-    return [...usableMetadataKeys].map(key => `${key}: ${this.metadata[key]}`).join(this.metadataSeparator);
-  }
-
-  getNodeInfo() {
-    return { start: this.startCharIdx, end: this.endCharIdx };
-  }
-
-  getText() {
-    return this.text;
-  }
-}
-
-/**
- * A document is just a special text node with a docId.
- */
-export class Document<T extends Metadata = Metadata> extends TextNode<T> {
-  constructor(init?: TextNodeParams<T>) {
-    super(init);
-  }
-
-  get type() {
-    return ObjectType.DOCUMENT;
-  }
-}
-
-function lazyInitHash(
-  value: ClassAccessorDecoratorTarget<BaseNode, string>,
-  _context: ClassAccessorDecoratorContext,
-): ClassAccessorDecoratorResult<BaseNode, string> {
-  return {
-    get() {
-      const oldValue = value.get.call(this);
-      if (oldValue === '') {
-        const hash = this.generateHash();
-        value.set.call(this, hash);
-      }
-      return value.get.call(this);
-    },
-    set(newValue: string) {
-      value.set.call(this, newValue);
-    },
-    init(value: string): string {
-      return value;
-    },
-  };
-}
-
-function createSHA256() {
-  const hash = createHash('sha256');
-  return {
-    update(data: string | Uint8Array): void {
-      hash.update(data);
-    },
-    digest() {
-      return hash.digest('base64');
-    },
-  };
-}

package/src/document/schema/types.ts

@@ -1,40 +0,0 @@
-export enum NodeRelationship {
-  SOURCE = 'SOURCE',
-  PREVIOUS = 'PREVIOUS',
-  NEXT = 'NEXT',
-  PARENT = 'PARENT',
-  CHILD = 'CHILD',
-}
-
-export enum ObjectType {
-  TEXT = 'TEXT',
-  IMAGE = 'IMAGE',
-  INDEX = 'INDEX',
-  DOCUMENT = 'DOCUMENT',
-  IMAGE_DOCUMENT = 'IMAGE_DOCUMENT',
-}
-
-export type Metadata = Record<string, any>;
-
-export interface RelatedNodeInfo<T extends Metadata = Metadata> {
-  nodeId: string;
-  nodeType?: ObjectType;
-  metadata: T;
-  hash?: string;
-}
-
-export type RelatedNodeType<T extends Metadata = Metadata> = RelatedNodeInfo<T> | RelatedNodeInfo<T>[];
-
-export type BaseNodeParams<T extends Metadata = Metadata> = {
-  id_?: string | undefined;
-  metadata?: T | undefined;
-  relationships?: Partial<Record<NodeRelationship, RelatedNodeType<T>>> | undefined;
-  hash?: string | undefined;
-};
-
-export type TextNodeParams<T extends Metadata = Metadata> = BaseNodeParams<T> & {
-  text?: string | undefined;
-  startCharIdx?: number | undefined;
-  endCharIdx?: number | undefined;
-  metadataSeparator?: string | undefined;
-};