Skip to Content
ReferenceRAGGraphRAG

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 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 Parameters

dimension?:

number
= 1536
Dimension of the embedding vectors

threshold?:

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

Methods

createGraph

Creates a knowledge graph from document chunks and their embeddings.

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

Parameters

chunks:

GraphChunk[]
Array of document chunks with text and metadata

embeddings:

GraphEmbedding[]
Array of embeddings corresponding to chunks

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[]

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

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 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
});
Metadata FiltersAstraVector