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:

1. Sets up a Mastra agent with gpt-4o-mini for response generation
2. Creates a vector query tool to manage vector store interactions
3. Uses existing embeddings to retrieve relevant context
4. 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:

OPENAI_API_KEY=your_openai_api_key_here
POSTGRES_CONNECTION_STRING=your_connection_string_here

Dependencies

Import the necessary dependencies:

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:

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:

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:

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

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 source on GitHub