@mastra/rag 0.0.2-alpha.41 → 0.0.2-alpha.43

package/CHANGELOG.md

@@ -1,5 +1,21 @@
 # @mastra/rag
 
+## 0.0.2-alpha.43
+
+### Patch Changes
+
+- Updated dependencies [b524c22]
+  - @mastra/core@0.1.27-alpha.59
+
+## 0.0.2-alpha.42
+
+### Patch Changes
+
+- 1874f40: Added re ranking tool to RAG
+- Updated dependencies [1874f40]
+- Updated dependencies [4b1ce2c]
+  - @mastra/core@0.1.27-alpha.58
+
 ## 0.0.2-alpha.41
 
 ### Patch Changes

package/dist/rag.cjs.development.js

@@ -3396,6 +3396,140 @@ var embed = function embed(chunk, options) {
   return core.embed(value, options);
 };
 
+// Default weights for different scoring components
+var DEFAULT_WEIGHTS = {
+  semantic: 0.4,
+  vector: 0.4,
+  position: 0.2
+};
+// Takes in a list of results from a vector store and reranks them based on semantic, vector, and position scores
+var RagReranker = /*#__PURE__*/function () {
+  function RagReranker(options) {
+    this.semanticProvider = void 0;
+    this.weights = void 0;
+    // Set up different weights for scoring components. Uses default weights if not provided
+    this.weights = _extends({}, DEFAULT_WEIGHTS, options.weights);
+    // Initialize semantic provider
+    if (options.semanticProvider === 'cohere') {
+      var _options$cohereModel;
+      if (!options.cohereApiKey) {
+        throw new Error('Cohere API key required when using Cohere provider');
+      }
+      this.semanticProvider = new core.CohereRelevanceScorer(options.cohereApiKey, (_options$cohereModel = options.cohereModel) != null ? _options$cohereModel : '');
+    } else {
+      if (!options.agentProvider) {
+        throw new Error('Agent provider options required when using Agent provider');
+      }
+      this.semanticProvider = new core.MastraAgentRelevanceScorer(options.agentProvider.provider, options.agentProvider.name);
+    }
+  }
+  // Calculate position score based on position in original list
+  var _proto = RagReranker.prototype;
+  _proto.calculatePositionScore = function calculatePositionScore(position, totalChunks) {
+    return 1 - position / totalChunks;
+  }
+  // Analyze query embedding features if needed
+  ;
+  _proto.analyzeQueryEmbedding = function analyzeQueryEmbedding(embedding) {
+    // Calculate embedding magnitude
+    var magnitude = Math.sqrt(embedding.reduce(function (sum, val) {
+      return sum + val * val;
+    }, 0));
+    // Find dominant features (highest absolute values)
+    var dominantFeatures = embedding.map(function (value, index) {
+      return {
+        value: Math.abs(value),
+        index: index
+      };
+    }).sort(function (a, b) {
+      return b.value - a.value;
+    }).slice(0, 5).map(function (item) {
+      return item.index;
+    });
+    return {
+      magnitude: magnitude,
+      dominantFeatures: dominantFeatures
+    };
+  }
+  // Adjust scores based on query characteristics
+  ;
+  _proto.adjustScores = function adjustScores(score, queryAnalysis) {
+    var magnitudeAdjustment = queryAnalysis.magnitude > 10 ? 1.1 : 1;
+    var featureStrengthAdjustment = queryAnalysis.magnitude > 5 ? 1.05 : 1;
+    return score * magnitudeAdjustment * featureStrengthAdjustment;
+  };
+  _proto.rerank = /*#__PURE__*/function () {
+    var _rerank = /*#__PURE__*/_asyncToGenerator(/*#__PURE__*/_regeneratorRuntime().mark(function _callee2(_ref) {
+      var _this = this;
+      var query, vectorStoreResults, queryEmbedding, _ref$topK, topK, resultLength, queryAnalysis, scoredResults;
+      return _regeneratorRuntime().wrap(function _callee2$(_context2) {
+        while (1) switch (_context2.prev = _context2.next) {
+          case 0:
+            query = _ref.query, vectorStoreResults = _ref.vectorStoreResults, queryEmbedding = _ref.queryEmbedding, _ref$topK = _ref.topK, topK = _ref$topK === void 0 ? 3 : _ref$topK;
+            resultLength = vectorStoreResults.length;
+            queryAnalysis = queryEmbedding ? this.analyzeQueryEmbedding(queryEmbedding) : null; // Get scores for each result
+            _context2.next = 5;
+            return Promise.all(vectorStoreResults.map(/*#__PURE__*/function () {
+              var _ref2 = _asyncToGenerator(/*#__PURE__*/_regeneratorRuntime().mark(function _callee(result, index) {
+                var _result$metadata;
+                var semanticScore, vectorScore, positionScore, finalScore;
+                return _regeneratorRuntime().wrap(function _callee$(_context) {
+                  while (1) switch (_context.prev = _context.next) {
+                    case 0:
+                      _context.next = 2;
+                      return _this.semanticProvider.getRelevanceScore(query, result == null || (_result$metadata = result.metadata) == null ? void 0 : _result$metadata.text);
+                    case 2:
+                      semanticScore = _context.sent;
+                      // Get existing vector score from result
+                      vectorScore = result.score; // Get score of vector based on position in original list
+                      positionScore = _this.calculatePositionScore(index, resultLength); // Combine scores using weights for each component
+                      finalScore = _this.weights.semantic * semanticScore + _this.weights.vector * vectorScore + _this.weights.position * positionScore;
+                      if (queryAnalysis) {
+                        finalScore = _this.adjustScores(finalScore, queryAnalysis);
+                      }
+                      return _context.abrupt("return", {
+                        result: result,
+                        score: finalScore,
+                        details: _extends({
+                          semantic: semanticScore,
+                          vector: vectorScore,
+                          position: positionScore
+                        }, queryAnalysis && {
+                          queryAnalysis: {
+                            magnitude: queryAnalysis.magnitude,
+                            dominantFeatures: queryAnalysis.dominantFeatures
+                          }
+                        })
+                      });
+                    case 8:
+                    case "end":
+                      return _context.stop();
+                  }
+                }, _callee);
+              }));
+              return function (_x2, _x3) {
+                return _ref2.apply(this, arguments);
+              };
+            }()));
+          case 5:
+            scoredResults = _context2.sent;
+            return _context2.abrupt("return", scoredResults.sort(function (a, b) {
+              return b.score - a.score;
+            }).slice(0, topK));
+          case 7:
+          case "end":
+            return _context2.stop();
+        }
+      }, _callee2, this);
+    }));
+    function rerank(_x) {
+      return _rerank.apply(this, arguments);
+    }
+    return rerank;
+  }();
+  return RagReranker;
+}();
+
 var createFilter = function createFilter(filter, vectorFilterType) {
   if (['pg', 'astra', 'pinecone'].includes(vectorFilterType)) {
     var _filter$keyword, _ref;
@@ -3418,14 +3552,44 @@ var createFilter = function createFilter(filter, vectorFilterType) {
     };
   }
 };
-var createVectorQueryTool = function createVectorQueryTool(_ref3) {
-  var vectorStoreName = _ref3.vectorStoreName,
-    indexName = _ref3.indexName,
-    _ref3$topK = _ref3.topK,
-    topK = _ref3$topK === void 0 ? 10 : _ref3$topK,
-    options = _ref3.options,
-    _ref3$vectorFilterTyp = _ref3.vectorFilterType,
-    vectorFilterType = _ref3$vectorFilterTyp === void 0 ? '' : _ref3$vectorFilterTyp;
+// Separate function to handle vector query search
+// Can be imported and used in custom tools
+var vectorQuerySearch = /*#__PURE__*/function () {
+  var _ref4 = /*#__PURE__*/_asyncToGenerator(/*#__PURE__*/_regeneratorRuntime().mark(function _callee(_ref3) {
+    var indexName, vectorStore, queryText, options, _ref3$queryFilter, queryFilter, topK, _yield$embed, embedding, results;
+    return _regeneratorRuntime().wrap(function _callee$(_context) {
+      while (1) switch (_context.prev = _context.next) {
+        case 0:
+          indexName = _ref3.indexName, vectorStore = _ref3.vectorStore, queryText = _ref3.queryText, options = _ref3.options, _ref3$queryFilter = _ref3.queryFilter, queryFilter = _ref3$queryFilter === void 0 ? {} : _ref3$queryFilter, topK = _ref3.topK;
+          _context.next = 3;
+          return embed(queryText, options);
+        case 3:
+          _yield$embed = _context.sent;
+          embedding = _yield$embed.embedding;
+          _context.next = 7;
+          return vectorStore.query(indexName, embedding, topK, queryFilter);
+        case 7:
+          results = _context.sent;
+          return _context.abrupt("return", results);
+        case 9:
+        case "end":
+          return _context.stop();
+      }
+    }, _callee);
+  }));
+  return function vectorQuerySearch(_x) {
+    return _ref4.apply(this, arguments);
+  };
+}();
+var createVectorQueryTool = function createVectorQueryTool(_ref5) {
+  var vectorStoreName = _ref5.vectorStoreName,
+    indexName = _ref5.indexName,
+    _ref5$topK = _ref5.topK,
+    topK = _ref5$topK === void 0 ? 10 : _ref5$topK,
+    options = _ref5.options,
+    _ref5$vectorFilterTyp = _ref5.vectorFilterType,
+    vectorFilterType = _ref5$vectorFilterTyp === void 0 ? '' : _ref5$vectorFilterTyp,
+    rerankOptions = _ref5.rerankOptions;
   return core.createTool({
     id: "VectorQuery " + vectorStoreName + " " + indexName + " Tool",
     inputSchema: zod.z.object({
@@ -3441,82 +3605,107 @@ var createVectorQueryTool = function createVectorQueryTool(_ref3) {
     }),
     description: "Fetches and combines the top " + topK + " relevant chunks from the " + vectorStoreName + " vector store using the " + indexName + " index",
     execute: function () {
-      var _execute = _asyncToGenerator(/*#__PURE__*/_regeneratorRuntime().mark(function _callee(_ref4) {
+      var _execute = _asyncToGenerator(/*#__PURE__*/_regeneratorRuntime().mark(function _callee2(_ref6) {
         var _mastra$vectors;
-        var _ref4$context, queryText, filter, mastra, relevantContext, vectorStore, _yield$embed, embedding, queryFilter, results, relevantChunks;
-        return _regeneratorRuntime().wrap(function _callee$(_context) {
-          while (1) switch (_context.prev = _context.next) {
+        var _ref6$context, queryText, filter, mastra, relevantContext, vectorStore, queryFilter, results, reranker, rerankedResults, _relevantChunks, relevantChunks;
+        return _regeneratorRuntime().wrap(function _callee2$(_context2) {
+          while (1) switch (_context2.prev = _context2.next) {
             case 0:
-              _ref4$context = _ref4.context, queryText = _ref4$context.queryText, filter = _ref4$context.filter, mastra = _ref4.mastra;
+              _ref6$context = _ref6.context, queryText = _ref6$context.queryText, filter = _ref6$context.filter, mastra = _ref6.mastra;
               relevantContext = '';
-              vectorStore = mastra == null || (_mastra$vectors = mastra.vectors) == null ? void 0 : _mastra$vectors[vectorStoreName];
-              _context.next = 5;
-              return embed(queryText, options);
-            case 5:
-              _yield$embed = _context.sent;
-              embedding = _yield$embed.embedding;
+              vectorStore = mastra == null || (_mastra$vectors = mastra.vectors) == null ? void 0 : _mastra$vectors[vectorStoreName]; // Get relevant chunks from the vector database
               if (!vectorStore) {
-                _context.next = 14;
+                _context2.next = 18;
                 break;
               }
               queryFilter = vectorFilterType && filter ? createFilter(filter, vectorFilterType) : {};
-              _context.next = 11;
-              return vectorStore.query(indexName, embedding, topK, queryFilter);
-            case 11:
-              results = _context.sent;
-              relevantChunks = results.map(function (result) {
+              _context2.next = 7;
+              return vectorQuerySearch({
+                indexName: indexName,
+                vectorStore: vectorStore,
+                queryText: queryText,
+                options: options,
+                queryFilter: queryFilter,
+                topK: topK
+              });
+            case 7:
+              results = _context2.sent;
+              if (!rerankOptions) {
+                _context2.next = 16;
+                break;
+              }
+              reranker = new RagReranker(rerankOptions);
+              _context2.next = 12;
+              return reranker.rerank({
+                query: queryText,
+                vectorStoreResults: results,
+                topK: topK
+              });
+            case 12:
+              rerankedResults = _context2.sent;
+              _relevantChunks = rerankedResults.map(function (_ref7) {
                 var _result$metadata;
+                var result = _ref7.result;
                 return result == null || (_result$metadata = result.metadata) == null ? void 0 : _result$metadata.text;
+              });
+              relevantContext = _relevantChunks.join('\n\n');
+              return _context2.abrupt("return", {
+                relevantContext: relevantContext
+              });
+            case 16:
+              relevantChunks = results.map(function (result) {
+                var _result$metadata2;
+                return result == null || (_result$metadata2 = result.metadata) == null ? void 0 : _result$metadata2.text;
               }); // Combine the chunks into a context string
               relevantContext = relevantChunks.join('\n\n');
-            case 14:
-              return _context.abrupt("return", {
+            case 18:
+              return _context2.abrupt("return", {
                 relevantContext: relevantContext
               });
-            case 15:
+            case 19:
             case "end":
-              return _context.stop();
+              return _context2.stop();
           }
-        }, _callee);
+        }, _callee2);
       }));
-      function execute(_x) {
+      function execute(_x2) {
         return _execute.apply(this, arguments);
       }
       return execute;
     }()
   });
 };
-var createDocumentChunker = function createDocumentChunker(_ref5) {
-  var doc = _ref5.doc,
-    _ref5$params = _ref5.params,
-    params = _ref5$params === void 0 ? {
+var createDocumentChunker = function createDocumentChunker(_ref8) {
+  var doc = _ref8.doc,
+    _ref8$params = _ref8.params,
+    params = _ref8$params === void 0 ? {
       strategy: 'recursive',
       size: 512,
       overlap: 50,
       separator: '\n'
-    } : _ref5$params;
+    } : _ref8$params;
   return core.createTool({
     id: "Document Chunker " + params.strategy + " " + params.size,
     inputSchema: zod.z.object({}),
     description: "Chunks document using " + params.strategy + " strategy with size " + params.size + " and " + params.overlap + " overlap",
     execute: function () {
-      var _execute2 = _asyncToGenerator(/*#__PURE__*/_regeneratorRuntime().mark(function _callee2() {
+      var _execute2 = _asyncToGenerator(/*#__PURE__*/_regeneratorRuntime().mark(function _callee3() {
         var chunks;
-        return _regeneratorRuntime().wrap(function _callee2$(_context2) {
-          while (1) switch (_context2.prev = _context2.next) {
+        return _regeneratorRuntime().wrap(function _callee3$(_context3) {
+          while (1) switch (_context3.prev = _context3.next) {
             case 0:
-              _context2.next = 2;
+              _context3.next = 2;
               return doc.chunk(params);
             case 2:
-              chunks = _context2.sent;
-              return _context2.abrupt("return", {
+              chunks = _context3.sent;
+              return _context3.abrupt("return", {
                 chunks: chunks
               });
             case 4:
             case "end":
-              return _context2.stop();
+              return _context3.stop();
           }
-        }, _callee2);
+        }, _callee3);
       }));
       function execute() {
         return _execute2.apply(this, arguments);
@@ -3531,8 +3720,10 @@ exports.MDocument = MDocument;
 exports.PgVector = PgVector;
 exports.PineconeVector = PineconeVector;
 exports.QdrantVector = QdrantVector;
+exports.RagReranker = RagReranker;
 exports.UpstashVector = UpstashVector;
 exports.createDocumentChunker = createDocumentChunker;
 exports.createVectorQueryTool = createVectorQueryTool;
 exports.embed = embed;
+exports.vectorQuerySearch = vectorQuerySearch;
 //# sourceMappingURL=rag.cjs.development.js.map