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(1536, 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

serialize

Returns a JSON-safe snapshot of the graph so it can be persisted and restored later instead of rebuilt with createGraph.

serialize(): GraphRAGSnapshot

Returns

Returns a GraphRAGSnapshot object containing:

version:

number

Snapshot format version, used to reject snapshots this version of the class can't load

dimension:

number

Dimension of the embedding vectors the graph was built with

threshold:

number

Similarity threshold the graph was built with

nodes:

GraphNode[]

All nodes in the graph, each including its full embedding

edges:

GraphEdge[]

All edges in the graph

The snapshot is a deep copy, so mutating it doesn't affect the graph it came from. Every node carries its full embedding, so snapshots are large: a 1,000-node graph built with 1536-dimension embeddings serializes to about 20 MB of JSON. Size your storage column accordingly.

deserialize

Rebuilds a GraphRAG instance from a snapshot produced by serialize.

static deserialize(snapshot: GraphRAGSnapshot): GraphRAG

Parameters

snapshot:

GraphRAGSnapshot

A snapshot previously returned by serialize

Throws if the snapshot version is unsupported, if a node embedding doesn't match the snapshot dimension, or if an edge references a node that isn't in the snapshot. A bad snapshot therefore fails at load time instead of during a later query.

Persisting a graph

Building a graph is O(n²) in the number of chunks, so rebuilding it on every process start is wasteful. Serialize the graph once and store the snapshot wherever you already keep state. A snapshot is plain JSON, so any store works (a file, a blob column, a key-value cache), and GraphRAG doesn't depend on a storage backend.

import { readFile, writeFile } from 'node:fs/promises'
import { GraphRAG } from '@mastra/rag'
import type { GraphRAGSnapshot } from '@mastra/rag'

const SNAPSHOT_PATH = './docs-graph.json'

async function loadOrBuildGraph() {
  try {
    const snapshot = JSON.parse(await readFile(SNAPSHOT_PATH, 'utf8')) as GraphRAGSnapshot
    return GraphRAG.deserialize(snapshot)
  } catch {
    // No usable snapshot yet, so build the graph from scratch
  }

  const graphRag = new GraphRAG(1536, 0.7)
  graphRag.createGraph(documentChunks, embeddings)

  await writeFile(SNAPSHOT_PATH, JSON.stringify(graphRag.serialize()))

  return graphRag
}

A snapshot reflects the chunks it was built from and isn't updated incrementally. When the underlying documents change, build the graph again and store a new snapshot.

Advanced example

// Stricter similarity threshold
const graphRag = new GraphRAG(1536, 0.8)

// 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,
})

Related

• createGraphRAGTool