@mastra/rag 2.7.3-alpha.1 → 2.8.0-alpha.3

package/dist/index.js

@@ -9,6 +9,7 @@ import { createSimilarityPrompt } from "@mastra/core/relevance";
 import ZeroEntropy from "zeroentropy";
 import { createTool } from "@mastra/core/tools";
 import { embedV1, embedV2, embedV3, embedV4 } from "@mastra/core/vector";
+import { getRequestUser } from "@mastra/core/utils";
 import { AgenticRetrieveStreamCommand, BedrockAgentRuntimeClient, RetrieveCommand } from "@aws-sdk/client-bedrock-agent-runtime";
 // \0rolldown/runtime.js
 var __create$1 = Object.create;
@@ -4847,6 +4848,7 @@ var RecursiveJsonTransformer = class RecursiveJsonTransformer {
 	ensureAscii;
 	convertLists;
 	constructor({ maxSize = 2e3, ensureAscii = false, convertLists = true }) {
+		if (!Number.isFinite(maxSize) || maxSize < 21) throw new RangeError("JSON maxSize must be at least 21");
 		this.maxSize = maxSize;
 		this.ensureAscii = ensureAscii;
 		this.convertLists = convertLists;
@@ -5779,6 +5781,7 @@ const validationSchemas = {
 		disallowedSpecial: setOrAllSchema
 	}).strict(),
 	json: baseChunkOptionsSchema.extend({
+		maxSize: z.number().min(21).optional(),
 		minSize: z.number().positive().optional(),
 		ensureAscii: z.boolean().optional(),
 		convertLists: z.boolean().optional()
@@ -6208,9 +6211,17 @@ async function executeRerank({ results, query, scorer, options }) {
 		...DEFAULT_WEIGHTS,
 		...options.weights
 	};
+	for (const [name, weight] of Object.entries(weights)) if (!Number.isFinite(weight) || weight < 0 || weight > 1) {
+		const err = /* @__PURE__ */ new Error(`Weights must be finite numbers between 0 and 1. Got ${name}: ${weight}`);
+		rerankSpan?.error({
+			error: err,
+			endSpan: true
+		});
+		throw err;
+	}
 	const sum = Object.values(weights).reduce((acc, w) => acc.plus(w.toString()), new Big(0));
 	if (!sum.eq(1)) {
-		const err = /* @__PURE__ */ new Error(`Weights must add up to 1. Got ${sum} from ${weights}`);
+		const err = /* @__PURE__ */ new Error(`Weights must add up to 1. Got ${sum} from ${JSON.stringify(weights)}`);
 		rerankSpan?.error({
 			error: err,
 			endSpan: true
@@ -6296,12 +6307,17 @@ var GraphRAG = class GraphRAG {
 	addNode(node) {
 		if (!node.embedding) throw new Error("Node must have an embedding");
 		if (node.embedding.length !== this.dimension) throw new Error(`Embedding dimension must be ${this.dimension}`);
-		this.nodes.set(node.id, node);
+		for (const component of node.embedding) if (!Number.isFinite(component)) throw new Error("Embedding must contain only finite numbers");
+		this.nodes.set(node.id, {
+			...node,
+			embedding: [...node.embedding]
+		});
 	}
 	addEdge(edge) {
 		if (!this.nodes.has(edge.source) || !this.nodes.has(edge.target)) throw new Error("Both source and target nodes must exist");
+		if (!Number.isFinite(edge.weight)) throw new Error("Edge weight must be a finite number");
 		this.edges.push(edge);
-		this.edges.push({
+		if (edge.source !== edge.target) this.edges.push({
 			source: edge.target,
 			target: edge.source,
 			weight: edge.weight,
@@ -6309,7 +6325,10 @@ var GraphRAG = class GraphRAG {
 		});
 	}
 	getNodes() {
-		return Array.from(this.nodes.values());
+		return Array.from(this.nodes.values(), (node) => ({
+			...node,
+			...node.embedding ? { embedding: [...node.embedding] } : {}
+		}));
 	}
 	getEdges() {
 		return this.edges;
@@ -6342,20 +6361,25 @@ var GraphRAG = class GraphRAG {
 	/**
 	* Rebuild a GraphRAG instance from a snapshot produced by `serialize()`.
 	*
-	* @throws if the snapshot version is unsupported, a node embedding does not
-	* match the snapshot dimension, or an edge references an unknown node.
+	* @throws if the snapshot version is unsupported, a node ID is duplicated,
+	* a node embedding does not match the snapshot dimension or contains nonfinite values,
+	* or an edge references an unknown node.
 	*/
 	static deserialize(snapshot) {
 		if (snapshot?.version !== GRAPH_RAG_SNAPSHOT_VERSION) throw new Error(`Unsupported GraphRAG snapshot version: ${snapshot?.version}`);
 		const graph = new GraphRAG(snapshot.dimension, snapshot.threshold);
-		for (const node of snapshot.nodes ?? []) graph.addNode({
-			...node,
-			...node.embedding ? { embedding: [...node.embedding] } : {},
-			...node.metadata ? { metadata: structuredClone(node.metadata) } : {}
-		});
+		for (const node of snapshot.nodes ?? []) {
+			if (graph.nodes.has(node.id)) throw new Error(`Duplicate node id in snapshot: ${node.id}`);
+			graph.addNode({
+				...node,
+				...node.embedding ? { embedding: [...node.embedding] } : {},
+				...node.metadata ? { metadata: structuredClone(node.metadata) } : {}
+			});
+		}
 		for (const edge of snapshot.edges ?? []) {
 			if (!graph.nodes.has(edge.source)) throw new Error(`Edge references unknown node: ${edge.source}`);
 			if (!graph.nodes.has(edge.target)) throw new Error(`Edge references unknown node: ${edge.target}`);
+			if (!Number.isFinite(edge.weight)) throw new Error("Edge weight must be a finite number");
 		}
 		graph.edges = (snapshot.edges ?? []).map((edge) => ({ ...edge }));
 		return graph;
@@ -6397,22 +6421,21 @@ var GraphRAG = class GraphRAG {
 	createGraph(chunks, embeddings) {
 		if (!chunks?.length || !embeddings?.length) throw new Error("Chunks and embeddings arrays must not be empty");
 		if (chunks.length !== embeddings.length) throw new Error("Chunks and embeddings must have the same length");
+		const graph = new GraphRAG(this.dimension, this.threshold);
 		chunks.forEach((chunk, index) => {
-			const node = {
+			graph.addNode({
 				id: index.toString(),
 				content: chunk.text,
 				embedding: embeddings[index]?.vector,
 				metadata: { ...chunk.metadata }
-			};
-			this.addNode(node);
-			this.nodes.set(node.id, node);
+			});
 		});
 		for (let i = 0; i < chunks.length; i++) {
 			const firstEmbedding = embeddings[i]?.vector;
 			for (let j = i + 1; j < chunks.length; j++) {
 				const secondEmbedding = embeddings[j]?.vector;
 				const similarity = this.cosineSimilarity(firstEmbedding, secondEmbedding);
-				if (similarity > this.threshold) this.addEdge({
+				if (similarity > this.threshold) graph.addEdge({
 					source: i.toString(),
 					target: j.toString(),
 					weight: similarity,
@@ -6420,6 +6443,8 @@ var GraphRAG = class GraphRAG {
 				});
 			}
 		}
+		this.nodes = graph.nodes;
+		this.edges = graph.edges;
 	}
 	selectWeightedNeighbor(neighbors) {
 		const totalWeight = neighbors.reduce((sum, n) => sum + n.weight, 0);
@@ -6463,9 +6488,10 @@ var GraphRAG = class GraphRAG {
 	*/
 	query({ query, topK = 10, randomWalkSteps = 100, restartProb = .15, filter }) {
 		if (!query || query.length !== this.dimension) throw new Error(`Query embedding must have dimension ${this.dimension}`);
-		if (topK < 1) throw new Error("TopK must be greater than 0");
-		if (randomWalkSteps < 1) throw new Error("Random walk steps must be greater than 0");
-		if (restartProb <= 0 || restartProb >= 1) throw new Error("Restart probability must be between 0 and 1");
+		for (const component of query) if (!Number.isFinite(component)) throw new Error("Query embedding must contain only finite numbers");
+		if (!Number.isInteger(topK) || topK < 1) throw new Error("TopK must be a positive integer");
+		if (!Number.isInteger(randomWalkSteps) || randomWalkSteps < 1) throw new Error("Random walk steps must be a positive integer");
+		if (!Number.isFinite(restartProb) || restartProb <= 0 || restartProb >= 1) throw new Error("Restart probability must be between 0 and 1");
 		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);
@@ -6723,7 +6749,8 @@ async function resolveVectorStore(options, context, logger) {
 		if (typeof vectorStoreOption === "function") {
 			const resolved = await vectorStoreOption({
 				requestContext,
-				mastra
+				mastra,
+				user: getRequestUser(requestContext)
 			});
 			if (!isValidMastraVector(resolved)) {
 				const contextStr = buildContextString(context);
@@ -6827,8 +6854,6 @@ const createGraphRAGTool = (options) => {
 		...defaultGraphOptions,
 		...options.graphOptions || {}
 	};
-	const graphRag = new GraphRAG(graphOptions.dimension, graphOptions.threshold);
-	let isInitialized = false;
 	const inputSchema = options.enableFilter ? filterSchema : z.object(baseSchema).passthrough();
 	return createTool({
 		id: toolId,
@@ -6888,38 +6913,40 @@ const createGraphRAGTool = (options) => {
 					observabilityContext
 				});
 				if (logger) logger.debug("vectorQuerySearch returned results", { count: results.length });
-				if (!isInitialized) {
-					const chunks = results.map((result) => ({
-						text: result?.metadata?.text,
-						metadata: result.metadata ?? {}
-					}));
-					const embeddings = results.map((result) => ({ vector: result.vector || [] }));
-					if (logger) logger.debug("Initializing graph", {
-						chunkCount: chunks.length,
-						embeddingCount: embeddings.length
-					});
-					const buildSpan = parentSpan?.createChildSpan({
-						type: SpanType.GRAPH_ACTION,
-						name: "graph build",
-						input: { nodeCount: chunks.length },
-						attributes: {
-							action: "build",
-							nodeCount: chunks.length,
-							threshold: graphOptions.threshold
-						}
-					});
-					try {
-						graphRag.createGraph(chunks, embeddings);
-					} catch (err) {
-						buildSpan?.error({
-							error: err,
-							endSpan: true
-						});
-						throw err;
+				if (results.length === 0) return {
+					relevantContext: [],
+					sources: []
+				};
+				const graphRag = new GraphRAG(graphOptions.dimension, graphOptions.threshold);
+				const chunks = results.map((result) => ({
+					text: result?.metadata?.text,
+					metadata: result.metadata ?? {}
+				}));
+				const embeddings = results.map((result) => ({ vector: result.vector || [] }));
+				if (logger) logger.debug("Initializing graph", {
+					chunkCount: chunks.length,
+					embeddingCount: embeddings.length
+				});
+				const buildSpan = parentSpan?.createChildSpan({
+					type: SpanType.GRAPH_ACTION,
+					name: "graph build",
+					input: { nodeCount: chunks.length },
+					attributes: {
+						action: "build",
+						nodeCount: chunks.length,
+						threshold: graphOptions.threshold
 					}
-					buildSpan?.end();
-					isInitialized = true;
-				} else if (logger) logger.debug("Graph already initialized, skipping graph construction");
+				});
+				try {
+					graphRag.createGraph(chunks, embeddings);
+				} catch (err) {
+					buildSpan?.error({
+						error: err,
+						endSpan: true
+					});
+					throw err;
+				}
+				buildSpan?.end();
 				const traverseSpan = parentSpan?.createChildSpan({
 					type: SpanType.GRAPH_ACTION,
 					name: "graph traverse",