@mastra/rag 2.1.3 → 2.2.0

package/dist/index.js

@@ -1,3 +1,4 @@
+import { SpanType, createObservabilityContext } from '@mastra/core/observability';
 import { Agent, isSupportedLanguageModel } from '@mastra/core/agent';
 import { randomUUID, createHash } from 'crypto';
 import { z } from 'zod';
@@ -6911,15 +6912,47 @@ var MDocument = class _MDocument {
     const textSplit = rt.transformDocuments(this.chunks);
     this.chunks = textSplit;
   }
-  async chunk(params) {
+  async chunk(params, options) {
     const { strategy: passedStrategy, extract, ...chunkOptions } = params || {};
     const strategy = passedStrategy || this.defaultStrategy();
-    validateChunkParams(strategy, chunkOptions);
-    await this.chunkBy(strategy, chunkOptions);
-    if (extract) {
-      await this.extractMetadata(extract);
+    const parentSpan = options?.observabilityContext?.tracingContext?.currentSpan;
+    const chunkSpan = parentSpan?.createChildSpan({
+      type: SpanType.RAG_ACTION,
+      name: `rag chunk: ${strategy}`,
+      input: { strategy },
+      attributes: {
+        action: "chunk",
+        strategy,
+        chunkSize: chunkOptions?.size,
+        chunkOverlap: chunkOptions?.overlap
+      }
+    });
+    try {
+      validateChunkParams(strategy, chunkOptions);
+      await this.chunkBy(strategy, chunkOptions);
+      if (extract) {
+        const extractSpan = chunkSpan?.createChildSpan({
+          type: SpanType.RAG_ACTION,
+          name: "rag extract metadata",
+          attributes: {
+            action: "extract_metadata",
+            extractor: Object.keys(extract).join(",")
+          }
+        });
+        try {
+          await this.extractMetadata(extract);
+        } catch (err) {
+          extractSpan?.error({ error: err, endSpan: true });
+          throw err;
+        }
+        extractSpan?.end();
+      }
+      chunkSpan?.end({ output: { chunkCount: this.chunks.length } });
+      return this.chunks;
+    } catch (err) {
+      chunkSpan?.error({ error: err, endSpan: true });
+      throw err;
     }
-    return this.chunks;
   }
   getDocs() {
     return this.chunks;
@@ -7041,14 +7074,27 @@ async function executeRerank({
   scorer,
   options
 }) {
-  const { queryEmbedding, topK = 3 } = options;
+  const { queryEmbedding, topK = 3, observabilityContext } = options;
+  const rerankSpan = observabilityContext?.tracingContext?.currentSpan?.createChildSpan({
+    type: SpanType.RAG_ACTION,
+    name: `rag rerank`,
+    input: { query, candidateCount: results.length },
+    attributes: {
+      action: "rerank",
+      candidateCount: results.length,
+      topN: topK,
+      scorer: scorer?.constructor?.name
+    }
+  });
   const weights = {
     ...DEFAULT_WEIGHTS,
     ...options.weights
   };
   const sum = Object.values(weights).reduce((acc, w) => acc.plus(w.toString()), new Big(0));
   if (!sum.eq(1)) {
-    throw new Error(`Weights must add up to 1. Got ${sum} from ${weights}`);
+    const err = new Error(`Weights must add up to 1. Got ${sum} from ${weights}`);
+    rerankSpan?.error({ error: err, endSpan: true });
+    throw err;
   }
   const resultLength = results.length;
   const queryAnalysis = queryEmbedding ? analyzeQueryEmbedding(queryEmbedding) : null;
@@ -7081,7 +7127,9 @@ async function executeRerank({
       };
     })
   );
-  return scoredResults.sort((a, b) => b.score - a.score).slice(0, topK);
+  const final = scoredResults.sort((a, b) => b.score - a.score).slice(0, topK);
+  rerankSpan?.end({ output: { returned: final.length } });
+  return final;
 }
 async function rerankWithScorer({
   results,
@@ -7371,33 +7419,67 @@ var vectorQuerySearch = async ({
   includeVectors = false,
   maxRetries = 2,
   databaseConfig = {},
-  providerOptions
+  providerOptions,
+  observabilityContext
 }) => {
+  const parentSpan = observabilityContext?.tracingContext?.currentSpan;
+  const embedSpan = parentSpan?.createChildSpan({
+    type: SpanType.RAG_EMBEDDING,
+    name: `rag embed: query`,
+    input: queryText,
+    attributes: {
+      mode: "query",
+      model: model?.modelId,
+      provider: model?.provider,
+      inputCount: 1
+    }
+  });
   let embeddingResult;
-  if (model.specificationVersion === "v3") {
-    embeddingResult = await embedV3({
-      model,
-      value: queryText,
-      maxRetries,
-      // Type assertion needed: providerOptions type is a union, but embedV3 expects specific version
-      ...providerOptions && { providerOptions }
-    });
-  } else if (model.specificationVersion === "v2") {
-    embeddingResult = await embedV2({
-      model,
-      value: queryText,
-      maxRetries,
-      // Type assertion needed: providerOptions type is a union, but embedV2 expects specific version
-      ...providerOptions && { providerOptions }
-    });
-  } else {
-    embeddingResult = await embedV1({
-      value: queryText,
-      model,
-      maxRetries
-    });
+  try {
+    if (model.specificationVersion === "v3") {
+      embeddingResult = await embedV3({
+        model,
+        value: queryText,
+        maxRetries,
+        // Type assertion needed: providerOptions type is a union, but embedV3 expects specific version
+        ...providerOptions && {
+          providerOptions
+        }
+      });
+    } else if (model.specificationVersion === "v2") {
+      embeddingResult = await embedV2({
+        model,
+        value: queryText,
+        maxRetries,
+        // Type assertion needed: providerOptions type is a union, but embedV2 expects specific version
+        ...providerOptions && {
+          providerOptions
+        }
+      });
+    } else {
+      embeddingResult = await embedV1({
+        value: queryText,
+        model,
+        maxRetries
+      });
+    }
+  } catch (err) {
+    embedSpan?.error({ error: err, endSpan: true });
+    throw err;
   }
   const embedding = embeddingResult.embedding;
+  const embedUsage = embeddingResult?.usage;
+  embedSpan?.end({
+    attributes: {
+      dimensions: embedding?.length,
+      ...embedUsage && {
+        usage: {
+          inputTokens: embedUsage.tokens ?? embedUsage.promptTokens ?? embedUsage.inputTokens
+        }
+      }
+    },
+    output: { dimensions: embedding?.length }
+  });
   const queryParams = {
     indexName,
     queryVector: embedding,
@@ -7405,7 +7487,29 @@ var vectorQuerySearch = async ({
     filter: queryFilter,
     includeVector: includeVectors
   };
-  const results = await vectorStore.query({ ...queryParams, ...databaseSpecificParams(databaseConfig) });
+  const querySpan = parentSpan?.createChildSpan({
+    type: SpanType.RAG_VECTOR_OPERATION,
+    name: `rag vector: query`,
+    // Pass filter as-is; the observability layer's deepClean handles
+    // size limits and sanitization centrally.
+    input: { topK, filter: queryFilter },
+    attributes: {
+      operation: "query",
+      indexName,
+      topK,
+      dimensions: embedding?.length
+    }
+  });
+  let results;
+  try {
+    results = await vectorStore.query({ ...queryParams, ...databaseSpecificParams(databaseConfig) });
+  } catch (err) {
+    querySpan?.error({ error: err, endSpan: true });
+    throw err;
+  }
+  querySpan?.end({
+    output: { returned: results?.length ?? 0 }
+  });
   return { results, queryEmbedding: embedding };
 };
 var databaseSpecificParams = (databaseConfig) => {
@@ -7669,7 +7773,9 @@ var createGraphRAGTool = (options) => {
     outputSchema,
     description: toolDescription,
     execute: async (inputData, context) => {
-      const { requestContext, mastra } = context || {};
+      const { requestContext, mastra, tracingContext } = context || {};
+      const observabilityContext = createObservabilityContext(tracingContext);
+      const parentSpan = observabilityContext.tracingContext?.currentSpan;
       const indexName = requestContext?.get("indexName") ?? options.indexName;
       const vectorStoreName = "vectorStore" in options ? storeName : requestContext?.get("vectorStoreName") ?? storeName;
       if (!indexName) throw new Error(`indexName is required, got: ${indexName}`);
@@ -7707,7 +7813,8 @@ var createGraphRAGTool = (options) => {
           queryFilter: Object.keys(queryFilter || {}).length > 0 ? queryFilter : void 0,
           topK: topKValue,
           includeVectors: true,
-          providerOptions
+          providerOptions,
+          observabilityContext
         });
         if (logger) {
           logger.debug("vectorQuerySearch returned results", { count: results.length });
@@ -7723,17 +7830,50 @@ var createGraphRAGTool = (options) => {
           if (logger) {
             logger.debug("Initializing graph", { chunkCount: chunks.length, embeddingCount: embeddings.length });
           }
-          graphRag.createGraph(chunks, embeddings);
+          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;
+          }
+          buildSpan?.end();
           isInitialized = true;
         } else if (logger) {
           logger.debug("Graph already initialized, skipping graph construction");
         }
-        const rerankedResults = graphRag.query({
-          query: queryEmbedding,
-          topK: topKValue,
-          randomWalkSteps,
-          restartProb
+        const traverseSpan = parentSpan?.createChildSpan({
+          type: SpanType.GRAPH_ACTION,
+          name: "graph traverse",
+          input: { topK: topKValue, randomWalkSteps, restartProb },
+          attributes: {
+            action: "traverse",
+            startNodes: 1,
+            maxDepth: randomWalkSteps
+          }
         });
+        let rerankedResults;
+        try {
+          rerankedResults = graphRag.query({
+            query: queryEmbedding,
+            topK: topKValue,
+            randomWalkSteps,
+            restartProb
+          });
+        } catch (err) {
+          traverseSpan?.error({ error: err, endSpan: true });
+          throw err;
+        }
+        traverseSpan?.end({ output: { returned: rerankedResults.length } });
         if (logger) {
           logger.debug("GraphRAG query returned results", { count: rerankedResults.length });
         }
@@ -7772,7 +7912,8 @@ var createVectorQueryTool = (options) => {
     inputSchema,
     outputSchema,
     execute: async (inputData, context) => {
-      const { requestContext, mastra } = context || {};
+      const { requestContext, mastra, tracingContext } = context || {};
+      const observabilityContext = createObservabilityContext(tracingContext);
       const indexName = requestContext?.get("indexName") ?? options.indexName;
       const vectorStoreName = "vectorStore" in options ? storeName : requestContext?.get("vectorStoreName") ?? storeName;
       const includeVectors = requestContext?.get("includeVectors") ?? options.includeVectors ?? false;
@@ -7813,7 +7954,8 @@ var createVectorQueryTool = (options) => {
           topK: topKValue,
           includeVectors,
           databaseConfig,
-          providerOptions
+          providerOptions,
+          observabilityContext
         });
         if (logger) {
           logger.debug("vectorQuerySearch returned results", { count: results.length });
@@ -7830,13 +7972,15 @@ var createVectorQueryTool = (options) => {
               scorer: reranker.model,
               options: {
                 ...reranker.options,
-                topK: reranker.options?.topK || topKValue
+                topK: reranker.options?.topK || topKValue,
+                observabilityContext
               }
             });
           } else {
             rerankedResults = await rerank(results, queryText, reranker.model, {
               ...reranker.options,
-              topK: reranker.options?.topK || topKValue
+              topK: reranker.options?.topK || topKValue,
+              observabilityContext
             });
           }
           if (logger) {