UNPKG

@mastra/rag

Version:

The Retrieval-Augmented Generation (RAG) module contains document processing and embedding utilities.

217 lines (167 loc) • 6.82 kB
> Guide on graph-based retrieval in Mastra # GraphRAG Graph-based retrieval enhances traditional vector search by following relationships between chunks of information. This approach is useful when information is spread across multiple documents or when documents reference each other. ## When to use GraphRAG GraphRAG is particularly effective when: - Information is spread across multiple documents - Documents reference each other - You need to traverse relationships to find complete answers - Understanding connections between concepts is important - Simple vector similarity misses important contextual relationships For straightforward semantic search without relationship traversal, use [standard retrieval methods](https://mastra.ai/docs/v1/rag/retrieval). ## How GraphRAG works GraphRAG combines vector similarity with knowledge graph traversal: 1. Initial vector search retrieves relevant chunks based on semantic similarity 2. A knowledge graph is constructed from the retrieved chunks 3. The graph is traversed to find connected information 4. Results include both directly relevant chunks and related content This process helps surface information that might not be semantically similar to the query but is contextually relevant through connections. ## Creating a graph query tool The Graph Query Tool provides agents with the ability to perform graph-based retrieval: ```ts import { createGraphRAGTool } from "@mastra/rag"; import { ModelRouterEmbeddingModel } from "@mastra/core/llm"; const graphQueryTool = createGraphRAGTool({ vectorStoreName: "pgVector", indexName: "embeddings", model: new ModelRouterEmbeddingModel("openai/text-embedding-3-small"), graphOptions: { threshold: 0.7, }, }); ``` ### Configuration options The `graphOptions` parameter controls how the knowledge graph is built and traversed: - `threshold`: Similarity threshold (0-1) for determining which chunks are related. Higher values create sparser graphs with stronger connections; lower values create denser graphs with more potential relationships. - `dimension`: Vector embedding dimension. Must match the embedding model's output dimension (e.g., 1536 for OpenAI's text-embedding-3-small). ```ts const graphQueryTool = createGraphRAGTool({ vectorStoreName: "pgVector", indexName: "embeddings", model: new ModelRouterEmbeddingModel("openai/text-embedding-3-small"), graphOptions: { dimension: 1536, threshold: 0.7, }, }); ``` ## Using GraphRAG with agents Integrate the graph query tool with an agent to enable graph-based retrieval: ```ts import { Agent } from "@mastra/core/agent"; const ragAgent = new Agent({ id: "rag-agent", name: "GraphRAG Agent", instructions: `You are a helpful assistant that answers questions based on the provided context. When answering questions, use the graph query tool to find relevant information and relationships. Base your answers on the context provided by the tool, and clearly state if the context doesn't contain enough information.`, model: "openai/gpt-5.1", tools: { graphQueryTool, }, }); ``` ## Document processing and storage Before using graph-based retrieval, process documents into chunks and store their embeddings: ```ts import { MDocument } from "@mastra/rag"; import { embedMany } from "ai"; import { ModelRouterEmbeddingModel } from "@mastra/core/llm"; // Create and chunk document const doc = MDocument.fromText("Your document content here..."); const chunks = await doc.chunk({ strategy: "recursive", size: 512, overlap: 50, separator: "\n", }); // Generate embeddings const { embeddings } = await embedMany({ model: new ModelRouterEmbeddingModel("openai/text-embedding-3-small"), values: chunks.map((chunk) => chunk.text), }); // Store in vector database const vectorStore = mastra.getVector("pgVector"); await vectorStore.createIndex({ indexName: "embeddings", dimension: 1536, }); await vectorStore.upsert({ indexName: "embeddings", vectors: embeddings, metadata: chunks?.map((chunk) => ({ text: chunk.text })), }); ``` ## Querying with GraphRAG Once configured, the agent can perform graph-based queries: ```ts const query = "What are the effects of infrastructure changes on local businesses?"; const response = await ragAgent.generate(query); console.log(response.text); ``` The agent uses the graph query tool to: 1. Convert the query to an embedding 2. Find semantically similar chunks in the vector store 3. Build a knowledge graph from related chunks 4. Traverse the graph to find connected information 5. Return comprehensive context for generating the response ## Choosing the right threshold The threshold parameter significantly impacts retrieval quality: - **High threshold (0.8-0.9)**: Strict connections, fewer relationships, more precise but potentially incomplete results - **Medium threshold (0.6-0.8)**: Balanced approach, good for most use cases - **Low threshold (0.4-0.6)**: More connections, broader context, risk of including less relevant information Start with 0.7 and adjust based on your specific use case: ```ts // Strict connections for precise answers const strictGraphTool = createGraphRAGTool({ vectorStoreName: "pgVector", indexName: "embeddings", model: new ModelRouterEmbeddingModel("openai/text-embedding-3-small"), graphOptions: { threshold: 0.85, }, }); // Broader connections for exploratory queries const broadGraphTool = createGraphRAGTool({ vectorStoreName: "pgVector", indexName: "embeddings", model: new ModelRouterEmbeddingModel("openai/text-embedding-3-small"), graphOptions: { threshold: 0.5, }, }); ``` ## Combining with other retrieval methods GraphRAG can be used alongside other retrieval approaches: ```ts import { createVectorQueryTool } from "@mastra/rag"; const vectorQueryTool = createVectorQueryTool({ vectorStoreName: "pgVector", indexName: "embeddings", model: new ModelRouterEmbeddingModel("openai/text-embedding-3-small"), }); const graphQueryTool = createGraphRAGTool({ vectorStoreName: "pgVector", indexName: "embeddings", model: new ModelRouterEmbeddingModel("openai/text-embedding-3-small"), graphOptions: { threshold: 0.7, }, }); const agent = new Agent({ id: "rag-agent", name: "RAG Agent", instructions: `Use vector search for simple fact-finding queries. Use graph search when you need to understand relationships or find connected information.`, model: "openai/gpt-5.1", tools: { vectorQueryTool, graphQueryTool, }, }); ``` This gives the agent flexibility to choose the appropriate retrieval method based on the query. ## Reference For detailed API documentation, see: - [GraphRAG Class](https://mastra.ai/reference/v1/rag/graph-rag) - [createGraphRAGTool()](https://mastra.ai/reference/v1/tools/graph-rag-tool)