Simple Retrieval Augmented Generation using Java

Last Updated : 20 Aug, 2025

Retrieval-Augmented Generation (RAG) is an innovative approach in the field of Natural Language Processing (NLP) that combines the strengths of retrieval-based and generation-based models to enhance the quality of generated text. In this article, we will define a complete working example of RAG.

Overview of the RAG Workflow

Load PDF Document
Generate Embedding (OpenAI)
Store in Vector DB (e.g., PostgreSQL + pgvector)
Accept Question via API
Embed Question
Retrieve Similar Documents
Generate Answer using OpenAI (GPT)
Return Answer

Prerequisites:
Download and Install Java Development Kit.
Download and Install IntelliJ.
Build Tool, ex. Gradle.

Step-by-step implementation of RAG

Step 1: Project Setup

The build.gradle file is the build configuration file used by Gradle (the build tool) to define dependencies and plugins for the project.

build.gradle:

plugins {
id 'java'
id 'org.springframework.boot' version '3.2.0'
}
repositories {
mavenCentral()
}
dependencies {
implementation 'org.springframework.boot:spring-boot-starter-web'
implementation 'dev.langchain4j:langchain4j-openai:0.24.0'
implementation 'org.apache.pdfbox:pdfbox:2.0.27'
implementation 'org.postgresql:postgresql:42.7.3'
implementation 'org.springframework.boot:spring-boot-starter-data-jdbc'
}

Step 2: Create DocumentLoader.java

The DocumentLoader.java class extracts text from PDF files using Apache PDFBox and generates embeddings via AiService. It then stores the text and embeddings into a vector store for semantic search and RAG-based answer generation.

DocumentLoader.java

Java

package com.example.rag;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.text.PDFTextStripper;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.io.File;
import java.util.UUID;

@Component
public class DocumentLoader {

    @Autowired
    private AiService aiService;

    @Autowired
    private VectorStore vectorStore;

    public void loadPdf(File pdfFile) throws Exception {
        PDDocument document = PDDocument.load(pdfFile);
        PDFTextStripper stripper = new PDFTextStripper();
        String text = stripper.getText(document);
        document.close();

        float[] embedding = aiService.embed(text);

        vectorStore.store(UUID.randomUUID().toString(), text, embedding);
    }
}

Explanation:

The class is marked with @Component so it can be auto-wired into other classes.
Loads and extracts text using Apache PDFBox.
Uses AIService to convert the text into a vector using OpenAI's embedding API.
Saves the embedding and text into a vector store (e.g, pgvector) using VectorStore.
Prepares documents for semantic search and answer generation.

Step 3: Create AiService.java class

The AiService.java class generates vector embeddings for text using OpenAI's embedding API. It also prepares prompts by combining documents and user questions for answer generation using a language model.

AiService.java:

Java

package com.example.rag;

import org.springframework.stereotype.Service;

import java.util.List;
import java.util.stream.Collectors;

@Service
public class AiService {

    public float[] embed(String text) {
        // TODO: Implement OpenAI embedding API call
        // For demo, returning dummy values
        return new float[]{0.1f, 0.2f, 0.3f};
    }

    public String generateAnswer(String question, List<String> docs) {
        String context = docs.stream().collect(Collectors.joining("\n"));
        String prompt = context + "\n\nQuestion: " + question;

        // TODO: Call OpenAI ChatCompletion API here
        return "Simulated GPT Answer for: " + question;
    }
}

Explanation:

@Service-annotated class is a Spring service component for business logic, allowing it to be auto-wired into other classes.
embed(String text) method to call OpenAI’s /v1/embeddings API, converting input text into a numerical vector (float[]).
embedding is used for tasks like semantic search, document similarity and vector storage.
generateAnswer(String question, List docs) method combines the list of documents and the user's question into a single prompt.

Step 4: Create VectorStore.java interface

The VectorStore.java interface defines methods to store document embeddings and perform similarity-based search. It enables retrieval of semantically relevant documents for RAG processing.

VectorStore.java

Java

public interface VectorStore {
    void store(String id, String content,
               float[] embedding); // Save document
               
    vector List<String> search(float[] queryEmbedding); // Search for similar vectors
                                    
}

Explanation:

The store method saves text and its embedding to a vector database.
The search method retrieves documents most similar to a given embedding.
Enables vector-based retrieval for AI tasks like contextual search and RAG.

Step 5. Create PgVectorStore.java class

The PgVectorStore.java class implements the VectorStore interface to store and search document embeddings in a PostgreSQL database using the pgvector extension. It uses JdbcTemplate to insert embeddings and retrieve the most similar documents based on vector similarity.

PgVectorStore.java:

Java

package com.example.rag;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.stereotype.Repository;

import java.util.Arrays;
import java.util.List;
import java.util.stream.Collectors;

@Repository
public class PgVectorStore implements VectorStore {

    @Autowired
    private JdbcTemplate jdbc;

    @Override
    public void store(String id, String content, float[] embedding) {
        jdbc.update(
            "INSERT INTO documents (id, content, embedding) VALUES (?, ?, ?)",
            id, content, toPgVector(embedding)
        );
    }

    @Override
    public List<String> search(float[] queryEmbedding) {
        String sql = "SELECT content FROM documents ORDER BY embedding <-> ? LIMIT 3";
        return jdbc.query(
            sql,
            new Object[]{toPgVector(queryEmbedding)},
            (rs, rowNum) -> rs.getString("content")
        );
    }

    private String toPgVector(float[] vector) {
        return Arrays.stream(vector)
                .mapToObj(Float::toString)
                .collect(Collectors.joining(", ", "[", "]"));
    }
}

Explanation:

@Repository, annotated class, is a Spring-managed data access component.
It implements the VectorStore interface to store and search vectors in a PostgreSQL database using pgvector.
Uses JdbcTemplate () for executing SQL queries.
search method retrieves the top 3 most similar documents using pgvector’s <-> similarity operator.
The store method is used insert the document ID, content and vector into the documents table.
toPgVector() method converts a Java float[] into the pgvector format ([v1, v2, ...]) for database storage and comparison.

Step 6: Create RagService.java class

The RagService.java class handles the core RAG logic by retrieving relevant documents from the VectorStore using vector similarity. It then combines those documents with the user question and uses AiService to generate a context-aware answer

RagService.java:

Java

package com.example.rag;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.util.List;

@Service
public class RagService {

    @Autowired
    private AiService aiService;

    @Autowired
    private VectorStore vectorStore;

    public String getAnswer(String question) {
        float[] embedding = aiService.embed(question);
        List<String> docs = vectorStore.search(embedding);
        return aiService.generateAnswer(question, docs);
    }
}

Step 7: Create RagController.java class

The RagController.java class exposes a REST endpoint to receive user questions via HTTP POST requests. It delegates the question to RagService and returns the AI-generated answer as the HTTP response

RagController.java

Java

package com.example.rag;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.*;

@RestController
@RequestMapping("/api")
public class RagController {

    @Autowired
    private RagService ragService;

    @PostMapping("/chat")
    public ResponseEntity<String> ask(@RequestBody String question) {
        String answer = ragService.getAnswer(question);
        return ResponseEntity.ok(answer);
    }
}

Explanation:

Marked with @RestController, making it a Spring REST controller that handles HTTP requests.
Base path set to /api/rag using @RequestMapping.
Injects RagService using @Autowired to handle business logic.
Defines a POST endpoint at /ask that accepts a user question in the request body.
Calls ragService.getAnswer(question) to generate an answer using RAG logic.
Returns the answer wrapped in a ResponseEntity as the HTTP response

Step 8: Create Main Application RagApplication.java

Java

package com.example.rag;

import org.springframework.boot.SpringApplication;
import org.springframework.boot.autoconfigure.SpringBootApplication;

@SpringBootApplication
public class RagApplication {
    public static void main(String[] args) {
        SpringApplication.run(RagApplication.class, args);
    }
}