Skip to main content

GraphRAG

The GraphRAG class implements a graph-based approach to retrieval augmented generation. It creates a knowledge graph from document chunks where nodes represent documents and edges represent semantic relationships, enabling both direct similarity matching and discovery of related content through graph traversal.

Basic UsageDirect link to Basic Usage

import { GraphRAG } from "@mastra/rag";

const graphRag = new GraphRAG({
  dimension: 1536,
  threshold: 0.7,
});

// Create the graph from chunks and embeddings
graphRag.createGraph(documentChunks, embeddings);

// Query the graph with embedding
const results = await graphRag.query({
  query: queryEmbedding,
  topK: 10,
  randomWalkSteps: 100,
  restartProb: 0.15,
});

Constructor ParametersDirect link to Constructor Parameters

dimension?:

number
= 1536
Dimension of the embedding vectors

threshold?:

number
= 0.7
Similarity threshold for creating edges between nodes (0-1)

MethodsDirect link to Methods

createGraphDirect link to createGraph

Creates a knowledge graph from document chunks and their embeddings.

createGraph(chunks: GraphChunk[], embeddings: GraphEmbedding[]): void

ParametersDirect link to Parameters

chunks:

GraphChunk[]
Array of document chunks with text and metadata

embeddings:

GraphEmbedding[]
Array of embeddings corresponding to chunks

queryDirect link to query

Performs a graph-based search combining vector similarity and graph traversal.

query({
  query,
  topK = 10,
  randomWalkSteps = 100,
  restartProb = 0.15
}: {
  query: number[];
  topK?: number;
  randomWalkSteps?: number;
  restartProb?: number;
}): RankedNode[]

ParametersDirect link to Parameters

query:

number[]
Query embedding vector

topK?:

number
= 10
Number of results to return

randomWalkSteps?:

number
= 100
Number of steps in random walk

restartProb?:

number
= 0.15
Probability of restarting walk from query node

ReturnsDirect link to Returns

Returns an array of RankedNode objects, where each node contains:

id:

string
Unique identifier for the node

content:

string
Text content of the document chunk

metadata:

Record<string, any>
Additional metadata associated with the chunk

score:

number
Combined relevance score from graph traversal

Advanced ExampleDirect link to Advanced Example

const graphRag = new GraphRAG({
  dimension: 1536,
  threshold: 0.8, // Stricter similarity threshold
});

// Create graph from chunks and embeddings
graphRag.createGraph(documentChunks, embeddings);

// Query with custom parameters
const results = await graphRag.query({
  query: queryEmbedding,
  topK: 5,
  randomWalkSteps: 200,
  restartProb: 0.2,
});