@mastra/rag 0.0.0-commonjs-20250227130920

package/src/graph-rag/index.test.ts

@@ -0,0 +1,235 @@
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+
+import type { GraphChunk, GraphEdge, GraphEmbedding, GraphNode } from './';
+import { GraphRAG } from './';
+
+describe('GraphRAG', () => {
+  beforeEach(() => {
+    vi.clearAllMocks(); // Clear any mock state before each test
+  });
+
+  describe('addNode', () => {
+    it('should throw an error if node does not have an embedding', () => {
+      const graph = new GraphRAG();
+      const node = {
+        id: '1',
+        content: 'Node 1',
+      };
+      expect(() => graph.addNode(node)).toThrow('Node must have an embedding');
+    });
+
+    it('should throw an error if node embedding dimension is not equal to the graph dimension', () => {
+      const graph = new GraphRAG(2);
+      const node: GraphNode = {
+        id: '1',
+        content: 'Node 1',
+        embedding: [1, 2, 3],
+      };
+      expect(() => graph.addNode(node)).toThrow('Embedding dimension must be 2');
+    });
+
+    it('should add a node to the graph', () => {
+      const graph = new GraphRAG(3);
+      const node = {
+        id: '1',
+        content: 'Node 1',
+        embedding: [1, 2, 3],
+      };
+      graph.addNode(node);
+      expect(graph['nodes'].size).toBe(1);
+    });
+  });
+
+  describe('addEdge', () => {
+    it('should throw an error if either source or target node does not exist', () => {
+      const graph = new GraphRAG();
+      const edge: GraphEdge = {
+        source: '1',
+        target: '2',
+        weight: 0.5,
+        type: 'semantic',
+      };
+      expect(() => graph.addEdge(edge)).toThrow('Both source and target nodes must exist');
+    });
+
+    it('should add an edge between two nodes', () => {
+      const graph = new GraphRAG(3);
+      const node1: GraphNode = {
+        id: '1',
+        content: 'Node 1',
+        embedding: [1, 2, 3],
+      };
+      const node2: GraphNode = {
+        id: '2',
+        content: 'Node 2',
+        embedding: [4, 5, 6],
+      };
+      graph.addNode(node1);
+      graph.addNode(node2);
+      const edge: GraphEdge = {
+        source: '1',
+        target: '2',
+        weight: 0.5,
+        type: 'semantic',
+      };
+      graph.addEdge(edge);
+      expect(graph['edges'].length).toBe(2);
+    });
+  });
+
+  describe('createGraph', () => {
+    it("chunks and embeddings can't be empty", () => {
+      const graph = new GraphRAG(3);
+      const chunks: GraphChunk[] = [];
+      const embeddings: GraphEmbedding[] = [];
+      expect(() => graph.createGraph(chunks, embeddings)).toThrowError(
+        'Chunks and embeddings arrays must not be empty',
+      );
+    });
+    it('chunks and embeddings must have the same length', () => {
+      const graph = new GraphRAG(3);
+      const chunks: GraphChunk[] = [
+        {
+          text: 'Chunk 1',
+          metadata: {},
+        },
+        {
+          text: 'Chunk 2',
+          metadata: {},
+        },
+      ];
+      const embeddings: GraphEmbedding[] = [
+        {
+          vector: [1, 2, 3],
+        },
+      ];
+      expect(() => graph.createGraph(chunks, embeddings)).toThrowError(
+        'Chunks and embeddings must have the same length',
+      );
+    });
+    it('should return the top ranked nodes', () => {
+      const results = [
+        {
+          metadata: {
+            text: 'Chunk 1',
+          },
+          vector: [1, 2, 3],
+        },
+        {
+          metadata: {
+            text: 'Chunk 2',
+          },
+          vector: [4, 5, 6],
+        },
+        {
+          metadata: {
+            text: 'Chunk 3',
+          },
+          vector: [7, 8, 9],
+        },
+      ];
+
+      const chunks = results.map(result => ({
+        text: result?.metadata?.text,
+        metadata: result.metadata,
+      }));
+      const embeddings = results.map(result => ({
+        vector: result.vector,
+      }));
+
+      const graph = new GraphRAG(3);
+      graph.createGraph(chunks, embeddings);
+
+      const nodes = graph.getNodes();
+      expect(nodes.length).toBe(3);
+      expect(nodes[0]?.id).toBe('0');
+      expect(nodes[1]?.id).toBe('1');
+      expect(nodes[2]?.id).toBe('2');
+
+      const edges = graph.getEdges();
+      expect(edges.length).toBe(6);
+    });
+  });
+
+  describe('query', () => {
+    it("query embedding can't be empty", () => {
+      const graph = new GraphRAG(3);
+      const queryEmbedding: number[] = [];
+      expect(() => graph.query({ query: queryEmbedding, topK: 2, randomWalkSteps: 3, restartProb: 0.1 })).toThrowError(
+        `Query embedding must have dimension ${3}`,
+      );
+    });
+
+    it('topK must be greater than 0', () => {
+      const graph = new GraphRAG(3);
+      const queryEmbedding = [1, 2, 3];
+      const topK = 0;
+      expect(() => graph.query({ query: queryEmbedding, topK, randomWalkSteps: 3, restartProb: 0.1 })).toThrowError(
+        'TopK must be greater than 0',
+      );
+    });
+
+    it('randomWalkSteps must be greater than 0', () => {
+      const graph = new GraphRAG(3);
+      const queryEmbedding = [1, 2, 3];
+      const topK = 2;
+      const randomWalkSteps = 0;
+      expect(() => graph.query({ query: queryEmbedding, topK, randomWalkSteps, restartProb: 0.1 })).toThrowError(
+        'Random walk steps must be greater than 0',
+      );
+    });
+
+    it('restartProb must be between 0 and 1', () => {
+      const graph = new GraphRAG(3);
+      const queryEmbedding = [1, 2, 3];
+      const topK = 2;
+      const randomWalkSteps = 3;
+      const restartProb = -0.1;
+      expect(() => graph.query({ query: queryEmbedding, topK, randomWalkSteps, restartProb })).toThrowError(
+        'Restart probability must be between 0 and 1',
+      );
+    });
+
+    it('should return the top ranked nodes', () => {
+      const graph = new GraphRAG(3);
+      const node1: GraphNode = {
+        id: '1',
+        content: 'Node 1',
+        embedding: [1, 2, 3],
+      };
+      const node2: GraphNode = {
+        id: '2',
+        content: 'Node 2',
+        embedding: [11, 12, 13],
+      };
+      const node3: GraphNode = {
+        id: '3',
+        content: 'Node 3',
+        embedding: [21, 22, 23],
+      };
+      graph.addNode(node1);
+      graph.addNode(node2);
+      graph.addNode(node3);
+      graph.addEdge({
+        source: '1',
+        target: '2',
+        weight: 0.5,
+        type: 'semantic',
+      });
+      graph.addEdge({
+        source: '2',
+        target: '3',
+        weight: 0.7,
+        type: 'semantic',
+      });
+
+      const queryEmbedding = [15, 16, 17];
+      const topK = 2;
+      const randomWalkSteps = 3;
+      const restartProb = 0.1;
+      const rerankedResults = graph.query({ query: queryEmbedding, topK, randomWalkSteps, restartProb });
+
+      expect(rerankedResults.length).toBe(2);
+    });
+  });
+});

package/src/graph-rag/index.ts

@@ -0,0 +1,306 @@
+/**
+ * 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>;
+}
+
+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

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

package/src/rerank/index.test.ts

@@ -0,0 +1,150 @@
+import { cohere } from '@ai-sdk/cohere';
+import { CohereRelevanceScorer } from '@mastra/core/relevance';
+import { describe, it, expect, vi, beforeEach } from 'vitest';
+
+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);
+  });
+});