Using the Vector Query Tool
This example demonstrates how to implement and use createVectorQueryTool
for semantic search in a RAG system. It shows how to configure the tool, manage vector storage, and retrieve relevant context effectively.
Overview
The system implements RAG using Mastra and OpenAI. Here’s what it does:
- Sets up a Mastra agent with gpt-4o-mini for response generation
- Creates a vector query tool to manage vector store interactions
- Uses existing embeddings to retrieve relevant context
- Generates context-aware responses using the Mastra agent
Note: To learn how to create and store embeddings, see the Upsert Embeddings guide.
Setup
Environment Setup
Make sure to set up your environment variables:
.env
OPENAI_API_KEY=your_openai_api_key_here
POSTGRES_CONNECTION_STRING=your_connection_string_here
Dependencies
Import the necessary dependencies:
src/index.ts
import { openai } from "@ai-sdk/openai";
import { Mastra } from "@mastra/core";
import { Agent } from "@mastra/core/agent";
import { createVectorQueryTool } from "@mastra/rag";
import { PgVector } from "@mastra/pg";
Vector Query Tool Creation
Create a tool that can query the vector database:
src/index.ts
const vectorQueryTool = createVectorQueryTool({
vectorStoreName: "pgVector",
indexName: "embeddings",
model: openai.embedding("text-embedding-3-small"),
});
Agent Configuration
Set up the Mastra agent that will handle the responses:
src/index.ts
export const ragAgent = new Agent({
name: "RAG Agent",
instructions:
"You are a helpful assistant that answers questions based on the provided context. Keep your answers concise and relevant.",
model: openai("gpt-4o-mini"),
tools: {
vectorQueryTool,
},
});
Instantiate PgVector and Mastra
Instantiate PgVector and Mastra with all components:
src/index.ts
const pgVector = new PgVector({ connectionString: process.env.POSTGRES_CONNECTION_STRING! });
export const mastra = new Mastra({
agents: { ragAgent },
vectors: { pgVector },
});
const agent = mastra.getAgent("ragAgent");
Example Usage
src/index.ts
const prompt = `
[Insert query based on document here]
Please base your answer only on the context provided in the tool.
If the context doesn't contain enough information to fully answer the question, please state that explicitly.
`;
const completion = await agent.generate(prompt);
console.log(completion.text);
View Example on GitHub