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: 5,
randomWalkSteps: 100,
restartProb: 0.15
});Constructor Parameters
dimension?:
number
Dimension of the embedding vectors
threshold?:
number
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[]): voidParameters
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: number[], topK?: number, randomWalkSteps?: number, restartProb?: number): RankedNode[]Parameters
query:
number[]
Query embedding vector
topK?:
number
Number of results to return
randomWalkSteps?:
number
Number of steps in random walk
restartProb?:
number
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
});