This project demonstrates a simple Retrieval Augmented Generation (RAG) pipeline using Spring AI, PGVector as the vector store, and OpenAI for embedding and chat models.
Or watch the tutorial on YouTube.
Components
- Spring Boot Application: The main application that orchestrates the RAG pipeline.
- PGVector: A PostgreSQL extension used as a vector database to store and retrieve document embeddings.
- OpenAI Embedding Model: Used to convert text documents and user queries into numerical vector representations.
- OpenAI Chat Model: Used to generate responses based on the user's query and retrieved relevant information.
- DocumentLoader: A Spring component that loads sample documents into the PGVector store on application startup.
- RagService: A service that handles the RAG logic:
- Takes a user query.
- Performs a similarity search in the PGVector store to find relevant documents.
- Constructs a prompt for the OpenAI chat model, incorporating the original query and the retrieved document content.
- Returns the generated response from the OpenAI chat model.
- RagController: A REST controller that exposes an endpoint (
/ai/rag) to interact with theRagService. rag-prompt.st: A prompt template used by theRagServiceto guide the AI model's response, ensuring it uses the provided information.
How to Run
Prerequisites
- Java 21: Ensure you have Java 21 installed.
- Gradle: This project uses Gradle for dependency management and building.
- PostgreSQL with PGVector Extension:
- Install PostgreSQL.
- Install the
pgvectorextension. You might need to compile it from source or use a pre-built image (e.g.,pgvector/pgvector). - Create a database (e.g.,
rag_demo).CREATE DATABASE rag_demo; - Enable the
vectorextension in your database (after connecting to your database\c rag_demo):CREATE EXTENSION IF NOT EXISTS vector; CREATE EXTENSION IF NOT EXISTS "uuid-ossp"; CREATE TABLE IF NOT EXISTS vector_store ( id uuid DEFAULT uuid_generate_v4() PRIMARY KEY, content text, embedding vector(768) ); CREATE INDEX ON vector_store USING HNSW (embedding vector_cosine_ops);
- OpenAI API Key: Obtain an API key from OpenAI.
Configuration
application.properties:- Update
src/main/resources/application.propertieswith your PostgreSQL connection details if they differ from the defaults. - Set your OpenAI API key as an environment variable named
OPENAI_API_KEY. Alternatively, you can directly paste it intoapplication.properties(not recommended for production):spring.ai.openai.api-key=YOUR_OPENAI_API_KEY
- Update
Build and Run
- Build the application:
./gradlew clean build
- Run the application:
The
OPENROUTER_API_KEY=YOUR_OPENROUTER_API_KEY ./gradlew bootRun
DocumentLoaderwill automatically load sample documents into your PGVector database on startup.
Test the RAG Endpoint
Once the application is running, you can access the RAG endpoint:
curl "http://localhost:8080/ai/rag?message=What is RAG?"You should receive a response generated by the OpenAI chat model, augmented with information retrieved from your PGVector store.