createGraphRAGTool()
The createGraphRAGTool()
creates a tool that enhances RAG by building a graph of semantic relationships between documents. It uses the GraphRAG
system under the hood to provide graph-based retrieval, finding relevant content through both direct similarity and connected relationships.
Usage Example
import { openai } from "@ai-sdk/openai";
import { createGraphRAGTool } from "@mastra/rag";
const graphTool = createGraphRAGTool({
vectorStoreName: "pinecone",
indexName: "docs",
model: openai.embedding("text-embedding-3-small"),
graphOptions: {
dimension: 1536,
threshold: 0.7,
randomWalkSteps: 100,
restartProb: 0.15,
},
});
Parameters
vectorStoreName:
string
Name of the vector store to query
indexName:
string
Name of the index within the vector store
model:
EmbeddingModel
Embedding model to use for vector search
enableFilter?:
boolean
= false
Enable filtering of results based on metadata
includeSources?:
boolean
= true
Include the full retrieval objects in the results
graphOptions?:
GraphOptions
= Default graph options
Configuration for the graph-based retrieval
description?:
string
Custom description for the tool. By default: 'Access and analyze relationships between information in the knowledge base to answer complex questions about connections and patterns'
GraphOptions
dimension?:
number
= 1536
Dimension of the embedding vectors
threshold?:
number
= 0.7
Similarity threshold for creating edges between nodes (0-1)
randomWalkSteps?:
number
= 100
Number of steps in random walk for graph traversal
restartProb?:
number
= 0.15
Probability of restarting random walk from query node
Returns
The tool returns an object with:
relevantContext:
string
Combined text from the most relevant document chunks, retrieved using graph-based ranking
sources:
QueryResult[]
Array of full retrieval result objects. Each object contains all information needed to reference the original document, chunk, and similarity score.
QueryResult object structure
{
id: string; // Unique chunk/document identifier
metadata: any; // All metadata fields (document ID, etc.)
vector: number[]; // Embedding vector (if available)
score: number; // Similarity score for this retrieval
document: string; // Full chunk/document text (if available)
}
Default Tool Description
The default description focuses on:
- Analyzing relationships between documents
- Finding patterns and connections
- Answering complex queries
Advanced Example
const graphTool = createGraphRAGTool({
vectorStoreName: "pinecone",
indexName: "docs",
model: openai.embedding("text-embedding-3-small"),
graphOptions: {
dimension: 1536,
threshold: 0.8, // Higher similarity threshold
randomWalkSteps: 200, // More exploration steps
restartProb: 0.2, // Higher restart probability
},
});
Example with Custom Description
const graphTool = createGraphRAGTool({
vectorStoreName: "pinecone",
indexName: "docs",
model: openai.embedding("text-embedding-3-small"),
description:
"Analyze document relationships to find complex patterns and connections in our company's historical data",
});
This example shows how to customize the tool description for a specific use case while maintaining its core purpose of relationship analysis.