@mastra/rag
Version:
The Retrieval-Augmented Generation (RAG) module contains document processing and embedding utilities.
189 lines (139 loc) • 5.91 kB
Markdown
> Guide on chunking and embedding documents in Mastra for efficient processing and retrieval.
# Chunking and Embedding Documents
Before processing, create a MDocument instance from your content. You can initialize it from various formats:
```ts
const docFromText = MDocument.fromText("Your plain text content...");
const docFromHTML = MDocument.fromHTML("<html>Your HTML content...</html>");
const docFromMarkdown = MDocument.fromMarkdown("# Your Markdown content...");
const docFromJSON = MDocument.fromJSON(`{ "key": "value" }`);
```
## Step 1: Document Processing
Use `chunk` to split documents into manageable pieces. Mastra supports multiple chunking strategies optimized for different document types:
- `recursive`: Smart splitting based on content structure
- `character`: Simple character-based splits
- `token`: Token-aware splitting
- `markdown`: Markdown-aware splitting
- `semantic-markdown`: Markdown splitting based on related header families
- `html`: HTML structure-aware splitting
- `json`: JSON structure-aware splitting
- `latex`: LaTeX structure-aware splitting
- `sentence`: Sentence-aware splitting
> **Note:**
Each strategy accepts different parameters optimized for its chunking approach.
Here's an example of how to use the `recursive` strategy:
```ts
const chunks = await doc.chunk({
strategy: "recursive",
maxSize: 512,
overlap: 50,
separators: ["\n"],
extract: {
metadata: true, // Optionally extract metadata
},
});
```
For text where preserving sentence structure is important, here's an example of how to use the `sentence` strategy:
```ts
const chunks = await doc.chunk({
strategy: "sentence",
maxSize: 450,
minSize: 50,
overlap: 0,
sentenceEnders: ["."],
});
```
For markdown documents where preserving the semantic relationships between sections is important, here's an example of how to use the `semantic-markdown` strategy:
```ts
const chunks = await doc.chunk({
strategy: "semantic-markdown",
joinThreshold: 500,
modelName: "gpt-3.5-turbo",
});
```
> **Note:**
Metadata extraction may use LLM calls, so ensure your API key is set.
We go deeper into chunking strategies in our [`chunk()` reference documentation](https://mastra.ai/reference/v1/rag/chunk).
## Step 2: Embedding Generation
Transform chunks into embeddings using your preferred provider. Mastra supports embedding models through the model router.
### Using the Model Router
The simplest way is to use Mastra's model router with `provider/model` strings:
```ts
import { ModelRouterEmbeddingModel } from "@mastra/core/llm";
import { embedMany } from "ai";
const { embeddings } = await embedMany({
model: new ModelRouterEmbeddingModel("openai/text-embedding-3-small"),
values: chunks.map((chunk) => chunk.text),
});
```
Mastra supports OpenAI and Google embedding models. For a complete list of supported embedding models, see the [embeddings reference](https://mastra.ai/reference/v1/rag/embeddings).
The model router automatically handles API key detection from environment variables.
The embedding functions return vectors, arrays of numbers representing the semantic meaning of your text, ready for similarity searches in your vector database.
### Configuring Embedding Dimensions
Embedding models typically output vectors with a fixed number of dimensions (e.g., 1536 for OpenAI's `text-embedding-3-small`).
Some models support reducing this dimensionality, which can help:
- Decrease storage requirements in vector databases
- Reduce computational costs for similarity searches
Here are some supported models:
OpenAI (text-embedding-3 models):
```ts
import { ModelRouterEmbeddingModel } from "@mastra/core/llm";
const { embeddings } = await embedMany({
model: new ModelRouterEmbeddingModel("openai/text-embedding-3-small"),
options: {
dimensions: 256, // Only supported in text-embedding-3 and later
},
values: chunks.map((chunk) => chunk.text),
});
```
Google (text-embedding-001):
```ts
const { embeddings } = await embedMany({
model: "google/gemini-embedding-001", {
outputDimensionality: 256, // Truncates excessive values from the end
}),
values: chunks.map((chunk) => chunk.text),
});
```
important[Vector Database Compatibility]
When storing embeddings, the vector database index must be configured to match the output size of your embedding model. If the dimensions do not match, you may get errors or data corruption.
## Example: Complete Pipeline
Here's an example showing document processing and embedding generation with both providers:
```ts
import { embedMany } from "ai";
import { MDocument } from "@mastra/rag";
// Initialize document
const doc = MDocument.fromText(`
Climate change poses significant challenges to global agriculture.
Rising temperatures and changing precipitation patterns affect crop yields.
`);
// Create chunks
const chunks = await doc.chunk({
strategy: "recursive",
maxSize: 256,
overlap: 50,
});
// Generate embeddings with OpenAI
import { ModelRouterEmbeddingModel } from "@mastra/core/llm";
const { embeddings } = await embedMany({
model: new ModelRouterEmbeddingModel("openai/text-embedding-3-small"),
values: chunks.map((chunk) => chunk.text),
});
// OR
// Generate embeddings with Cohere
const { embeddings } = await embedMany({
model: "cohere/embed-english-v3.0",
values: chunks.map((chunk) => chunk.text),
});
// Store embeddings in your vector database
await vectorStore.upsert({
indexName: "embeddings",
vectors: embeddings,
});
```
##
For more examples of different chunking strategies and embedding configurations, see:
- [Chunk Reference](https://mastra.ai/reference/v1/rag/chunk)
- [Embeddings Reference](https://mastra.ai/reference/v1/rag/embeddings)
For more details on vector databases and embeddings, see:
- [Vector Databases](./vector-databases)
- [Embedding API Reference](https://mastra.ai/reference/v1/rag/embeddings)