@mastra/rag 1.2.2 → 1.2.3-alpha.1

package/src/graph-rag/index.ts

@@ -1,306 +0,0 @@
-/**
- * TODO: GraphRAG Enhancements
- *  - Add support for more edge types (sequential, hierarchical, citation, etc)
- *  - Allow for custom edge types
- *  - Utilize metadata for richer connections
- *  - Improve graph traversal and querying using types
- */
-
-type SupportedEdgeType = 'semantic';
-
-// Types for graph nodes and edges
-export interface GraphNode {
-  id: string;
-  content: string;
-  embedding?: number[];
-  metadata?: Record<string, any>;
-}
-
-export interface RankedNode extends GraphNode {
-  score: number;
-}
-
-export interface GraphEdge {
-  source: string;
-  target: string;
-  weight: number;
-  type: SupportedEdgeType;
-}
-
-export interface GraphChunk {
-  text: string;
-  metadata: Record<string, any>;
-}
-
-export interface GraphEmbedding {
-  vector: number[];
-}
-
-export class GraphRAG {
-  private nodes: Map<string, GraphNode>;
-  private edges: GraphEdge[];
-  private dimension: number;
-  private threshold: number;
-
-  constructor(dimension: number = 1536, threshold: number = 0.7) {
-    this.nodes = new Map();
-    this.edges = [];
-    this.dimension = dimension;
-    this.threshold = threshold;
-  }
-
-  // Add a node to the graph
-  addNode(node: GraphNode): void {
-    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);
-  }
-
-  // Add an edge between two nodes
-  addEdge(edge: GraphEdge): void {
-    if (!this.nodes.has(edge.source) || !this.nodes.has(edge.target)) {
-      throw new Error('Both source and target nodes must exist');
-    }
-    this.edges.push(edge);
-    // Add reverse edge
-    this.edges.push({
-      source: edge.target,
-      target: edge.source,
-      weight: edge.weight,
-      type: edge.type,
-    });
-  }
-
-  // Helper method to get all nodes
-  getNodes(): GraphNode[] {
-    return Array.from(this.nodes.values());
-  }
-
-  // Helper method to get all edges
-  getEdges(): GraphEdge[] {
-    return this.edges;
-  }
-
-  getEdgesByType(type: string): GraphEdge[] {
-    return this.edges.filter(edge => edge.type === type);
-  }
-
-  clear(): void {
-    this.nodes.clear();
-    this.edges = [];
-  }
-
-  updateNodeContent(id: string, newContent: string): void {
-    const node = this.nodes.get(id);
-    if (!node) {
-      throw new Error(`Node ${id} not found`);
-    }
-    node.content = newContent;
-  }
-
-  // Get neighbors of a node
-  private getNeighbors(nodeId: string, edgeType?: string): { id: string; weight: number }[] {
-    return this.edges
-      .filter(edge => edge.source === nodeId && (!edgeType || edge.type === edgeType))
-      .map(edge => ({
-        id: edge.target,
-        weight: edge.weight,
-      }))
-      .filter(node => node !== undefined);
-  }
-
-  // Calculate cosine similarity between two vectors
-  private cosineSimilarity(vec1: number[], vec2: number[]): number {
-    if (!vec1 || !vec2) {
-      throw new Error('Vectors must not be null or undefined');
-    }
-    const vectorLength = vec1.length;
-
-    if (vectorLength !== vec2.length) {
-      throw new Error(`Vector dimensions must match: vec1(${vec1.length}) !== vec2(${vec2.length})`);
-    }
-
-    let dotProduct = 0;
-    let normVec1 = 0;
-    let normVec2 = 0;
-
-    for (let i = 0; i < vectorLength; i++) {
-      const a = vec1[i]!; // Non-null assertion operator
-      const b = vec2[i]!;
-
-      dotProduct += a * b;
-      normVec1 += a * a;
-      normVec2 += b * b;
-    }
-    const magnitudeProduct = Math.sqrt(normVec1 * normVec2);
-
-    if (magnitudeProduct === 0) {
-      return 0;
-    }
-
-    const similarity = dotProduct / magnitudeProduct;
-    return Math.max(-1, Math.min(1, similarity));
-  }
-
-  createGraph(chunks: GraphChunk[], embeddings: GraphEmbedding[]) {
-    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');
-    }
-    // Create nodes from chunks
-    chunks.forEach((chunk, index) => {
-      const node: GraphNode = {
-        id: index.toString(),
-        content: chunk.text,
-        embedding: embeddings[index]?.vector,
-        metadata: { ...chunk.metadata },
-      };
-      this.addNode(node);
-      this.nodes.set(node.id, node);
-    });
-
-    // Create edges based on cosine similarity
-    for (let i = 0; i < chunks.length; i++) {
-      const firstEmbedding = embeddings[i]?.vector as number[];
-      for (let j = i + 1; j < chunks.length; j++) {
-        const secondEmbedding = embeddings[j]?.vector as number[];
-        const similarity = this.cosineSimilarity(firstEmbedding, secondEmbedding);
-
-        // Only create edges if similarity is above threshold
-        if (similarity > this.threshold) {
-          this.addEdge({
-            source: i.toString(),
-            target: j.toString(),
-            weight: similarity,
-            type: 'semantic',
-          });
-        }
-      }
-    }
-  }
-
-  private selectWeightedNeighbor(neighbors: Array<{ id: string; weight: number }>): string {
-    // Sum all weights to normalize probabilities
-    const totalWeight = neighbors.reduce((sum, n) => sum + n.weight, 0);
-
-    // Pick a random point in the total weight range
-    let remainingWeight = Math.random() * totalWeight;
-
-    // Subtract each weight from our random value until we go below 0
-    // Higher weights will make us go below 0 more often, making them more likely to be selected
-    for (const neighbor of neighbors) {
-      remainingWeight -= neighbor.weight;
-      if (remainingWeight <= 0) {
-        return neighbor.id;
-      }
-    }
-
-    return neighbors[neighbors.length - 1]?.id as string;
-  }
-
-  // Perform random walk with restart
-  private randomWalkWithRestart(startNodeId: string, steps: number, restartProb: number): Map<string, number> {
-    const visits = new Map<string, number>();
-    let currentNodeId = startNodeId;
-
-    for (let step = 0; step < steps; step++) {
-      // Record visit
-      visits.set(currentNodeId, (visits.get(currentNodeId) || 0) + 1);
-
-      // Decide whether to restart
-      if (Math.random() < restartProb) {
-        currentNodeId = startNodeId;
-        continue;
-      }
-
-      // Get neighbors
-      const neighbors = this.getNeighbors(currentNodeId);
-      if (neighbors.length === 0) {
-        currentNodeId = startNodeId;
-        continue;
-      }
-
-      // Select random weighted neighbor and set as current node
-      currentNodeId = this.selectWeightedNeighbor(neighbors);
-    }
-
-    // Normalize visits
-    const totalVisits = Array.from(visits.values()).reduce((a, b) => a + b, 0);
-    const normalizedVisits = new Map<string, number>();
-    for (const [nodeId, count] of visits) {
-      normalizedVisits.set(nodeId, count / totalVisits);
-    }
-
-    return normalizedVisits;
-  }
-
-  // Retrieve relevant nodes using hybrid approach
-  query({
-    query,
-    topK = 10,
-    randomWalkSteps = 100,
-    restartProb = 0.15,
-  }: {
-    query: number[];
-    topK?: number;
-    randomWalkSteps?: number;
-    restartProb?: number;
-  }): RankedNode[] {
-    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');
-    }
-    // Retrieve nodes and calculate similarity
-    const similarities = Array.from(this.nodes.values()).map(node => ({
-      node,
-      similarity: this.cosineSimilarity(query, node.embedding!),
-    }));
-
-    // Sort by similarity
-    similarities.sort((a, b) => b.similarity - a.similarity);
-    const topNodes = similarities.slice(0, topK);
-
-    // Re-ranks nodes using random walk with restart
-    const rerankedNodes = new Map<string, { node: GraphNode; score: number }>();
-
-    // For each top node, perform random walk
-    for (const { node, similarity } of topNodes) {
-      const walkScores = this.randomWalkWithRestart(node.id, randomWalkSteps, restartProb);
-
-      // Combine dense retrieval score with graph score
-      for (const [nodeId, walkScore] of walkScores) {
-        const node = this.nodes.get(nodeId)!;
-        const existingScore = rerankedNodes.get(nodeId)?.score || 0;
-        rerankedNodes.set(nodeId, {
-          node,
-          score: existingScore + similarity * walkScore,
-        });
-      }
-    }
-
-    // Sort by final score and return top K nodes
-    return Array.from(rerankedNodes.values())
-      .sort((a, b) => b.score - a.score)
-      .slice(0, topK)
-      .map(item => ({
-        id: item.node.id,
-        content: item.node.content,
-        metadata: item.node.metadata,
-        score: item.score,
-      }));
-  }
-}

package/src/index.ts

@@ -1,8 +0,0 @@
-export * from './document/document';
-export * from './document/types';
-export * from './rerank';
-export * from './rerank/relevance';
-export { GraphRAG } from './graph-rag';
-export * from './tools';
-export * from './utils/vector-prompts';
-export * from './utils/default-settings';

package/src/rerank/index.test.ts

@@ -1,150 +0,0 @@
-import { cohere } from '@ai-sdk/cohere';
-import { describe, it, expect, vi, beforeEach } from 'vitest';
-import { CohereRelevanceScorer } from './relevance';
-
-import { rerank } from '.';
-
-vi.spyOn(CohereRelevanceScorer.prototype, 'getRelevanceScore').mockImplementation(async () => {
-  return 1;
-});
-
-const getScoreSpreads = (results1: any, results2: any) => {
-  const scoreSpread1 = Math.max(...results1.map((r: any) => r.score)) - Math.min(...results1.map((r: any) => r.score));
-  const scoreSpread2 = Math.max(...results2.map((r: any) => r.score)) - Math.min(...results2.map((r: any) => r.score));
-  return { scoreSpread1, scoreSpread2 };
-};
-
-describe('rerank', () => {
-  beforeEach(() => {
-    vi.clearAllMocks();
-  });
-
-  it('should throw an error if weights do not add up to 1', async () => {
-    const results = [
-      { id: '1', metadata: { text: 'Test result 1' }, score: 0.5 },
-      { id: '2', metadata: { text: 'Test result 2' }, score: 0.4 },
-      { id: '3', metadata: { text: 'Test result 3' }, score: 0.9 },
-    ];
-    await expect(
-      rerank(results, 'test query', cohere('rerank-v3.5'), { weights: { semantic: 0.5, vector: 0.3, position: 0.5 } }),
-    ).rejects.toThrow('Weights must add up to 1');
-  });
-
-  it('should rerank results with default weights', async () => {
-    const results = [
-      { id: '1', metadata: { text: 'Test result 1' }, score: 0.5 },
-      { id: '2', metadata: { text: 'Test result 2' }, score: 0.4 },
-      { id: '3', metadata: { text: 'Test result 3' }, score: 0.9 },
-    ];
-
-    const rerankedResults = await rerank(results, 'test query', cohere('rerank-v3.5'), { topK: 2 });
-
-    expect(rerankedResults).toHaveLength(2);
-    expect(rerankedResults[0]).toStrictEqual({
-      result: { id: '3', metadata: { text: 'Test result 3' }, score: 0.9 },
-      score: 0.8266666666666667,
-      details: {
-        semantic: 1,
-        vector: 0.9,
-        position: 0.33333333333333337,
-      },
-    });
-    expect(rerankedResults[1]).toStrictEqual({
-      result: { id: '1', metadata: { text: 'Test result 1' }, score: 0.5 },
-      score: 0.8,
-      details: {
-        semantic: 1,
-        vector: 0.5,
-        position: 1,
-      },
-    });
-
-    const { scoreSpread1, scoreSpread2 } = getScoreSpreads(results, rerankedResults);
-    expect(scoreSpread1).toBe(0.5);
-    expect(scoreSpread2).toBe(0.026666666666666616);
-  });
-
-  it('should rerank results with custom weights', async () => {
-    const results = [
-      { id: '1', metadata: { text: 'Test result 1' }, score: 0.5 },
-      { id: '2', metadata: { text: 'Test result 2' }, score: 0.4 },
-      { id: '3', metadata: { text: 'Test result 3' }, score: 0.9 },
-    ];
-
-    const rerankedResults = await rerank(results, 'test query', cohere('rerank-v3.5'), {
-      weights: {
-        semantic: 0.5,
-        vector: 0.4,
-        position: 0.1,
-      },
-      topK: 2,
-    });
-
-    expect(rerankedResults).toHaveLength(2);
-    expect(rerankedResults[0]).toStrictEqual({
-      result: { id: '3', metadata: { text: 'Test result 3' }, score: 0.9 },
-      score: 0.8933333333333334,
-      details: {
-        semantic: 1,
-        vector: 0.9,
-        position: 0.33333333333333337,
-      },
-    });
-    expect(rerankedResults[1]).toStrictEqual({
-      result: { id: '1', metadata: { text: 'Test result 1' }, score: 0.5 },
-      score: 0.7999999999999999,
-      details: {
-        semantic: 1,
-        vector: 0.5,
-        position: 1,
-      },
-    });
-    const { scoreSpread1, scoreSpread2 } = getScoreSpreads(results, rerankedResults);
-    expect(scoreSpread1).toBe(0.5);
-    expect(scoreSpread2).toBe(0.09333333333333349);
-  });
-
-  it('should handle query embedding when provided', async () => {
-    const results = [
-      { id: '1', metadata: { text: 'Test result 1' }, score: 0.8 },
-      { id: '2', metadata: { text: 'Test result 2' }, score: 0.6 },
-    ];
-
-    const rerankedResults = await rerank(results, 'test query', cohere('rerank-v3.5'), {
-      queryEmbedding: [0.5, 0.3, -0.2, 0.4],
-      topK: 2,
-    });
-
-    // Ensure query embedding analysis is being applied (we don't know exact score without knowing internals)
-    expect(rerankedResults).toHaveLength(2);
-    expect(rerankedResults[0]).toStrictEqual({
-      result: { id: '1', metadata: { text: 'Test result 1' }, score: 0.8 },
-      score: 0.9200000000000002,
-      details: {
-        semantic: 1,
-        vector: 0.8,
-        position: 1,
-        queryAnalysis: {
-          magnitude: 0.7348469228349535,
-          dominantFeatures: [0, 3, 1, 2],
-        },
-      },
-    });
-    expect(rerankedResults[1]).toStrictEqual({
-      result: { id: '2', metadata: { text: 'Test result 2' }, score: 0.6 },
-      score: 0.74,
-      details: {
-        semantic: 1,
-        vector: 0.6,
-        position: 0.5,
-        queryAnalysis: {
-          magnitude: 0.7348469228349535,
-          dominantFeatures: [0, 3, 1, 2],
-        },
-      },
-    });
-    const { scoreSpread1, scoreSpread2 } = getScoreSpreads(results, rerankedResults);
-    expect(scoreSpread1).toBe(0.20000000000000007);
-    expect(scoreSpread2).toBe(0.18000000000000016);
-  });
-});

package/src/rerank/index.ts

@@ -1,198 +0,0 @@
-import type { MastraLanguageModel } from '@mastra/core/agent';
-import type { RelevanceScoreProvider } from '@mastra/core/relevance';
-import type { QueryResult } from '@mastra/core/vector';
-import { Big } from 'big.js';
-import { MastraAgentRelevanceScorer, CohereRelevanceScorer } from './relevance';
-
-// Default weights for different scoring components (must add up to 1)
-const DEFAULT_WEIGHTS = {
-  semantic: 0.4,
-  vector: 0.4,
-  position: 0.2,
-} as const;
-
-type WeightConfig = {
-  semantic?: number;
-  vector?: number;
-  position?: number;
-};
-
-interface ScoringDetails {
-  semantic: number;
-  vector: number;
-  position: number;
-  queryAnalysis?: {
-    magnitude: number;
-    dominantFeatures: number[];
-  };
-}
-
-export interface RerankResult {
-  result: QueryResult;
-  score: number;
-  details: ScoringDetails;
-}
-
-// For use in the vector store tool
-export interface RerankerOptions {
-  weights?: WeightConfig;
-  topK?: number;
-}
-
-// For use in the rerank function
-export interface RerankerFunctionOptions {
-  weights?: WeightConfig;
-  queryEmbedding?: number[];
-  topK?: number;
-}
-
-export interface RerankConfig {
-  options?: RerankerOptions;
-  model: MastraLanguageModel | RelevanceScoreProvider;
-}
-
-// Calculate position score based on position in original list
-function calculatePositionScore(position: number, totalChunks: number): number {
-  return 1 - position / totalChunks;
-}
-
-// Analyze query embedding features if needed
-function analyzeQueryEmbedding(embedding: number[]): {
-  magnitude: number;
-  dominantFeatures: number[];
-} {
-  // Calculate embedding magnitude
-  const magnitude = Math.sqrt(embedding.reduce((sum, val) => sum + val * val, 0));
-
-  // Find dominant features (highest absolute values)
-  const dominantFeatures = embedding
-    .map((value, index) => ({ value: Math.abs(value), index }))
-    .sort((a, b) => b.value - a.value)
-    .slice(0, 5)
-    .map(item => item.index);
-
-  return { magnitude, dominantFeatures };
-}
-
-// Adjust scores based on query characteristics
-function adjustScores(score: number, queryAnalysis: { magnitude: number; dominantFeatures: number[] }): number {
-  const magnitudeAdjustment = queryAnalysis.magnitude > 10 ? 1.1 : 1;
-
-  const featureStrengthAdjustment = queryAnalysis.magnitude > 5 ? 1.05 : 1;
-
-  return score * magnitudeAdjustment * featureStrengthAdjustment;
-}
-
-async function executeRerank({
-  results,
-  query,
-  scorer,
-  options,
-}: {
-  results: QueryResult[];
-  query: string;
-  scorer: RelevanceScoreProvider;
-  options: RerankerFunctionOptions;
-}) {
-  const { queryEmbedding, topK = 3 } = options;
-  const weights = {
-    ...DEFAULT_WEIGHTS,
-    ...options.weights,
-  };
-
-  //weights must add up to 1
-  const sum = Object.values(weights).reduce((acc: Big, w: number) => 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 resultLength = results.length;
-
-  const queryAnalysis = queryEmbedding ? analyzeQueryEmbedding(queryEmbedding) : null;
-
-  // Get scores for each result
-  const scoredResults = await Promise.all(
-    results.map(async (result, index) => {
-      // Get semantic score from chosen provider
-      let semanticScore = 0;
-      if (result?.metadata?.text) {
-        semanticScore = await scorer.getRelevanceScore(query, result?.metadata?.text);
-      }
-
-      // Get existing vector score from result
-      const vectorScore = result.score;
-
-      // Get score of vector based on position in original list
-      const positionScore = calculatePositionScore(index, resultLength);
-
-      // Combine scores using weights for each component
-      let finalScore =
-        weights.semantic * semanticScore + weights.vector * vectorScore + weights.position * positionScore;
-
-      if (queryAnalysis) {
-        finalScore = adjustScores(finalScore, queryAnalysis);
-      }
-
-      return {
-        result,
-        score: finalScore,
-        details: {
-          semantic: semanticScore,
-          vector: vectorScore,
-          position: positionScore,
-          ...(queryAnalysis && {
-            queryAnalysis: {
-              magnitude: queryAnalysis.magnitude,
-              dominantFeatures: queryAnalysis.dominantFeatures,
-            },
-          }),
-        },
-      };
-    }),
-  );
-
-  // Sort by score and take top K
-  return scoredResults.sort((a, b) => b.score - a.score).slice(0, topK);
-}
-
-export async function rerankWithScorer({
-  results,
-  query,
-  scorer,
-  options,
-}: {
-  results: QueryResult[];
-  query: string;
-  scorer: RelevanceScoreProvider;
-  options: RerankerFunctionOptions;
-}): Promise<RerankResult[]> {
-  return executeRerank({
-    results,
-    query,
-    scorer,
-    options,
-  });
-}
-
-// Takes in a list of results from a vector store and reranks them based on semantic, vector, and position scores
-export async function rerank(
-  results: QueryResult[],
-  query: string,
-  model: MastraLanguageModel,
-  options: RerankerFunctionOptions,
-): Promise<RerankResult[]> {
-  let semanticProvider: RelevanceScoreProvider;
-
-  if (model.modelId === 'rerank-v3.5') {
-    semanticProvider = new CohereRelevanceScorer(model.modelId);
-  } else {
-    semanticProvider = new MastraAgentRelevanceScorer(model.provider, model);
-  }
-
-  return executeRerank({
-    results,
-    query,
-    scorer: semanticProvider,
-    options,
-  });
-}