@mastra/rag 1.3.5-alpha.0 → 1.3.6-alpha.0

package/dist/index.js

@@ -6808,7 +6808,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++) {
@@ -6817,7 +6817,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;
@@ -6831,12 +6834,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}`);
@@ -6850,15 +6863,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;