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
Parameter Requirements: Most fields can be set at creation as defaults.
Some fields can be overridden at runtime via the runtime context or input. If
a required field is missing from both creation and runtime, an error will be
thrown. Note that model
, id
, and description
can only be set at creation
time.
id?:
description?:
vectorStoreName:
indexName:
model:
enableFilter?:
includeSources?:
graphOptions?:
GraphOptions
dimension?:
threshold?:
randomWalkSteps?:
restartProb?:
Returns
The tool returns an object with:
relevantContext:
sources:
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.
Example: Using Runtime Context
const graphTool = createGraphRAGTool({
vectorStoreName: "pinecone",
indexName: "docs",
model: openai.embedding("text-embedding-3-small"),
});
When using runtime context, provide required parameters at execution time via the runtime context:
const runtimeContext = new RuntimeContext<{
vectorStoreName: string;
indexName: string;
topK: number;
filter: any;
}>();
runtimeContext.set("vectorStoreName", "my-store");
runtimeContext.set("indexName", "my-index");
runtimeContext.set("topK", 5);
runtimeContext.set("filter", { category: "docs" });
runtimeContext.set("randomWalkSteps", 100);
runtimeContext.set("restartProb", 0.15);
const response = await agent.generate(
"Find documentation from the knowledge base.",
{
runtimeContext,
},
);
For more information on runtime context, please see: