@mastra/rag 2.0.0-beta.1 → 2.0.0-beta.3

package/dist/index.js

@@ -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}"
@@ -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({
@@ -5561,6 +5585,11 @@ var MarkdownHeaderTransformer = class {
         currentContent.push(line);
         continue;
       }
+      const isTableLine = strippedLine.includes("|") && strippedLine.length > 0;
+      if (isTableLine) {
+        currentContent.push(line);
+        continue;
+      }
       let headerMatched = false;
       for (const [sep, name14] of this.headersToSplitOn) {
         if (strippedLine.startsWith(sep) && (strippedLine.length === sep.length || strippedLine[sep.length] === " ")) {
@@ -6381,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;
     }
@@ -6797,7 +6846,7 @@ var GraphRAG = class {
     return neighbors[neighbors.length - 1]?.id;
   }
   // Perform random walk with restart
-  randomWalkWithRestart(startNodeId, steps, restartProb) {
+  randomWalkWithRestart(startNodeId, steps, restartProb, allowedNodeIds) {
     const visits = /* @__PURE__ */ new Map();
     let currentNodeId = startNodeId;
     for (let step = 0; step < steps; step++) {
@@ -6806,7 +6855,10 @@ var GraphRAG = class {
         currentNodeId = startNodeId;
         continue;
       }
-      const neighbors = this.getNeighbors(currentNodeId);
+      let neighbors = this.getNeighbors(currentNodeId);
+      if (allowedNodeIds) {
+        neighbors = neighbors.filter((n) => allowedNodeIds.has(n.id));
+      }
       if (neighbors.length === 0) {
         currentNodeId = startNodeId;
         continue;
@@ -6820,12 +6872,22 @@ var GraphRAG = class {
     }
     return normalizedVisits;
   }
+  /**
+   * Query the graph with a dense embedding and optional metadata filter.
+   *
+   * @param query - The embedding vector to query.
+   * @param topK - Number of top results to return.
+   * @param randomWalkSteps - Steps for random walk reranking.
+   * @param restartProb - Restart probability for random walk.
+   * @param filter - Optional strict metadata filter. All key-value pairs must match exactly.
+   */
   // Retrieve relevant nodes using hybrid approach
   query({
     query,
     topK = 10,
     randomWalkSteps = 100,
-    restartProb = 0.15
+    restartProb = 0.15,
+    filter
   }) {
     if (!query || query.length !== this.dimension) {
       throw new Error(`Query embedding must have dimension ${this.dimension}`);
@@ -6839,15 +6901,20 @@ var GraphRAG = class {
     if (restartProb <= 0 || restartProb >= 1) {
       throw new Error("Restart probability must be between 0 and 1");
     }
-    const similarities = Array.from(this.nodes.values()).map((node) => ({
+    const filterEntries = Object.entries(filter ?? {});
+    const matchesFilter = (node) => filterEntries.length === 0 ? true : filterEntries.every(([key, value]) => node.metadata?.[key] === value);
+    const nodesToSearch = Array.from(this.nodes.values()).filter(matchesFilter);
+    const similarities = nodesToSearch.map((node) => ({
       node,
       similarity: this.cosineSimilarity(query, node.embedding)
     }));
     similarities.sort((a, b) => b.similarity - a.similarity);
     const topNodes = similarities.slice(0, topK);
+    const useFilter = filterEntries.length > 0;
+    const allowedNodeIds = useFilter ? new Set(nodesToSearch.map((n) => n.id)) : void 0;
     const rerankedNodes = /* @__PURE__ */ new Map();
     for (const { node, similarity } of topNodes) {
-      const walkScores = this.randomWalkWithRestart(node.id, randomWalkSteps, restartProb);
+      const walkScores = this.randomWalkWithRestart(node.id, randomWalkSteps, restartProb, allowedNodeIds);
       for (const [nodeId, walkScore] of walkScores) {
         const node2 = this.nodes.get(nodeId);
         const existingScore = rerankedNodes.get(nodeId)?.score || 0;
@@ -7130,9 +7197,9 @@ var createGraphRAGTool = (options) => {
               return typeof filter === "string" ? JSON.parse(filter) : filter;
             } catch (error) {
               if (logger) {
-                logger.warn("Failed to parse filter as JSON, using empty filter", { filter, error });
+                logger.error("Invalid filter", { filter, error });
               }
-              return {};
+              throw new Error(`Invalid filter format: ${error instanceof Error ? error.message : String(error)}`);
             }
           })();
         }
@@ -7257,9 +7324,9 @@ var createVectorQueryTool = (options) => {
               return typeof filter === "string" ? JSON.parse(filter) : filter;
             } catch (error) {
               if (logger) {
-                logger.warn("Failed to parse filter as JSON, using empty filter", { filter, error });
+                logger.error("Invalid filter", { filter, error });
               }
-              return {};
+              throw new Error(`Invalid filter format: ${error instanceof Error ? error.message : String(error)}`);
             }
           })();
         }