Skip to main content

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​

note

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?:

string
Custom ID for the tool. By default: 'GraphRAG {vectorStoreName} {indexName} Tool'. (Set at creation only.)

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.' (Set at creation only.)

vectorStoreName:

string
Name of the vector store to query. (Can be set at creation or overridden at runtime.)

indexName:

string
Name of the index within the vector store. (Can be set at creation or overridden at runtime.)

model:

EmbeddingModel
Embedding model to use for vector search. (Set at creation only.)

enableFilter?:

boolean
= false
Enable filtering of results based on metadata. (Set at creation only, but will be automatically enabled if a filter is provided in the runtime context.)

includeSources?:

boolean
= true
Include the full retrieval objects in the results. (Can be set at creation or overridden at runtime.)

graphOptions?:

GraphOptions
= Default graph options
Configuration for the graph-based retrieval

providerOptions?:

Record<string, Record<string, any>>
Provider-specific options for the embedding model (e.g., outputDimensionality). **Important**: Only works with AI SDK EmbeddingModelV2 models. For V1 models, configure options when creating the model itself.

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. (Can be set at creation or overridden at runtime.)

restartProb?:

number
= 0.15
Probability of restarting random walk from query node. (Can be set at creation or overridden at runtime.)

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.

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: