@mastra/rag 2.0.0-beta.2 → 2.0.0-beta.4

package/dist/index.js

@@ -1,4 +1,4 @@
-import { Agent } from '@mastra/core/agent';
+import { Agent, isSupportedLanguageModel } from '@mastra/core/agent';
 import { randomUUID, createHash } from 'crypto';
 import { z } from 'zod';
 import { parse } from 'node-html-better-parser';
@@ -7,7 +7,7 @@ import { Big } from 'big.js';
 import { createSimilarityPrompt } from '@mastra/core/relevance';
 import ZeroEntropy from 'zeroentropy';
 import { createTool } from '@mastra/core/tools';
-import { embedV2, embedV1 } from '@mastra/core/vector';
+import { embedV3, embedV2, embedV1 } from '@mastra/core/vector';
 
 var __create = Object.create;
 var __defProp = Object.defineProperty;
@@ -266,7 +266,7 @@ Provide keywords in the following comma-separated format: 'KEYWORDS: <keywords>'
 var 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. "
+  "Given the contextual information below, generate {numQuestions} questions this context can provide specific answers to which are unlikely to be found elsewhere. 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}"
@@ -4135,7 +4135,7 @@ var TitleExtractor = class extends BaseExtractor {
       const titleCandidates = await this.getTitlesCandidates(nodes);
       const combinedTitles = titleCandidates.join(", ");
       let title = "";
-      if (this.llm.specificationVersion === "v2") {
+      if (isSupportedLanguageModel(this.llm)) {
         const miniAgent = new Agent({
           id: "title-extractor",
           model: this.llm,
@@ -4174,7 +4174,7 @@ var TitleExtractor = class extends BaseExtractor {
     });
     const titleJobs = nodes.map(async (node) => {
       let completion;
-      if (this.llm.specificationVersion === "v2") {
+      if (isSupportedLanguageModel(this.llm)) {
         const result = await miniAgent.generate([
           { role: "user", content: this.nodeTemplate.format({ context: node.getContent() }) }
         ]);
@@ -4240,7 +4240,7 @@ var SummaryExtractor = class extends BaseExtractor {
       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") {
+    if (isSupportedLanguageModel(this.llm)) {
       const result = await miniAgent.generate([{ role: "user", content: prompt }]);
       summary = result.text;
     } else {
@@ -4326,7 +4326,7 @@ var QuestionsAnsweredExtractor = class extends BaseExtractor {
       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") {
+    if (isSupportedLanguageModel(this.llm)) {
       const result2 = await miniAgent.generate([{ role: "user", content: prompt }]);
       questionsText = result2.text;
     } else {
@@ -4398,7 +4398,7 @@ var KeywordExtractor = class extends BaseExtractor {
         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") {
+      if (isSupportedLanguageModel(this.llm)) {
         const result = await miniAgent.generate([
           {
             role: "user",
@@ -4993,23 +4993,47 @@ var HTMLSectionTransformer = class {
     }
     return "/" + parts.join("/");
   }
+  getTextContent(element) {
+    if (!element) return "";
+    if (!element.tagName) {
+      return element.text || "";
+    }
+    let content = element.text || "";
+    if (element.childNodes) {
+      for (const child of element.childNodes) {
+        const childText = this.getTextContent(child);
+        if (childText) {
+          content += " " + childText;
+        }
+      }
+    }
+    return content.trim();
+  }
   splitHtmlByHeaders(htmlDoc) {
     const sections = [];
     const root = parse(htmlDoc);
     const headers = Object.keys(this.headersToSplitOn);
     const headerElements = root.querySelectorAll(headers.join(","));
-    headerElements.forEach((headerElement, index) => {
+    headerElements.forEach((headerElement) => {
       const header = headerElement.text?.trim() || "";
       const tagName = headerElement.tagName;
       const xpath = this.getXPath(headerElement);
       let content = "";
-      let currentElement = headerElement.nextElementSibling;
-      const nextHeader = headerElements[index + 1];
-      while (currentElement && (!nextHeader || currentElement !== nextHeader)) {
-        if (currentElement.text) {
-          content += currentElement.text.trim() + " ";
+      const parentNode = headerElement.parentNode;
+      if (parentNode && parentNode.childNodes) {
+        let foundHeader = false;
+        for (const node of parentNode.childNodes) {
+          if (node === headerElement) {
+            foundHeader = true;
+            continue;
+          }
+          if (foundHeader && node.tagName && headers.includes(node.tagName.toLowerCase())) {
+            break;
+          }
+          if (foundHeader) {
+            content += this.getTextContent(node) + " ";
+          }
         }
-        currentElement = currentElement.nextElementSibling;
       }
       content = content.trim();
       sections.push({
@@ -6386,13 +6410,33 @@ var MDocument = class _MDocument {
   async chunkHTML(options) {
     if (options?.headers?.length) {
       const rt = new HTMLHeaderTransformer(options);
-      const textSplit = rt.transformDocuments(this.chunks);
+      let textSplit = rt.transformDocuments(this.chunks);
+      if (options?.maxSize) {
+        const textSplitter = new RecursiveCharacterTransformer({
+          maxSize: options.maxSize,
+          overlap: options.overlap,
+          keepSeparator: options.keepSeparator,
+          addStartIndex: options.addStartIndex,
+          stripWhitespace: options.stripWhitespace
+        });
+        textSplit = textSplitter.splitDocuments(textSplit);
+      }
       this.chunks = textSplit;
       return;
     }
     if (options?.sections?.length) {
       const rt = new HTMLSectionTransformer(options);
-      const textSplit = rt.transformDocuments(this.chunks);
+      let textSplit = rt.transformDocuments(this.chunks);
+      if (options?.maxSize) {
+        const textSplitter = new RecursiveCharacterTransformer({
+          maxSize: options.maxSize,
+          overlap: options.overlap,
+          keepSeparator: options.keepSeparator,
+          addStartIndex: options.addStartIndex,
+          stripWhitespace: options.stripWhitespace
+        });
+        textSplit = textSplitter.splitDocuments(textSplit);
+      }
       this.chunks = textSplit;
       return;
     }
@@ -6544,7 +6588,7 @@ Always return just the number, no explanation.`,
     const prompt = createSimilarityPrompt(query, text);
     const model = await this.agent.getModel();
     let response;
-    if (model.specificationVersion === "v2") {
+    if (isSupportedLanguageModel(model)) {
       response = await this.agent.generate(prompt);
     } else {
       response = await this.agent.generateLegacy(prompt);
@@ -6930,7 +6974,14 @@ var vectorQuerySearch = async ({
   providerOptions
 }) => {
   let embeddingResult;
-  if (model.specificationVersion === "v2") {
+  if (model.specificationVersion === "v3") {
+    embeddingResult = await embedV3({
+      model,
+      value: queryText,
+      maxRetries,
+      ...providerOptions && { providerOptions }
+    });
+  } else if (model.specificationVersion === "v2") {
     embeddingResult = await embedV2({
       model,
       value: queryText,