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

<PropertiesTable content={[ { name: 'dimension', type: 'number', description: 'Dimension of the embedding vectors', isOptional: true, defaultValue: '1536', }, { name: 'threshold', type: 'number', description: 'Similarity threshold for creating edges between nodes (0-1)', isOptional: true, defaultValue: '0.7', }, ]} />

Methods

createGraph

Creates a knowledge graph from document chunks and their embeddings.

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

Parameters

<PropertiesTable content={[ { name: 'chunks', type: 'GraphChunk[]', description: 'Array of document chunks with text and metadata', isOptional: false, }, { name: 'embeddings', type: 'GraphEmbedding[]', description: 'Array of embeddings corresponding to chunks', isOptional: false, }, ]} />

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

<PropertiesTable content={[ { name: 'query', type: 'number[]', description: 'Query embedding vector', isOptional: false, }, { name: 'topK', type: 'number', description: 'Number of results to return', isOptional: true, defaultValue: '10', }, { name: 'randomWalkSteps', type: 'number', description: 'Number of steps in random walk', isOptional: true, defaultValue: '100', }, { name: 'restartProb', type: 'number', description: 'Probability of restarting walk from query node', isOptional: true, defaultValue: '0.15', }, ]} />

Returns

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

<PropertiesTable content={[ { name: 'id', type: 'string', description: 'Unique identifier for the node', }, { name: 'content', type: 'string', description: 'Text content of the document chunk', }, { name: 'metadata', type: 'Record<string, any>', description: 'Additional metadata associated with the chunk', }, { name: 'score', type: 'number', description: '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,
})

Related

• createGraphRAGTool