Retrieval Augmented Generation (RAG)

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-starter-model-openai</artifactId>
</dependency>
<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-starter-vector-store-pgvector</artifactId>
</dependency>
<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-advisors-vector-store</artifactId>
</dependency>
<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-rag</artifactId>
</dependency>
<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-pdf-document-reader</artifactId>
</dependency>

@Service
public class PolicyIngestionService {

    private static final Logger log = LoggerFactory.getLogger(PolicyIngestionService.class);
    private final VectorStore vectorStore;

    public PolicyIngestionService(VectorStore vectorStore) {
        this.vectorStore = vectorStore;
    }

    public void ingest(Resource pdf, String docId) {
        // 1. Read: one Document per PDF page, with page number metadata
        List<Document> pages = new PagePdfDocumentReader(pdf).get();

        // 2. Transform: split pages into token-sized chunks
        List<Document> chunks = new TokenTextSplitter().apply(pages);
        chunks.forEach(c -> c.getMetadata().put("docId", docId));

        // 3. Write: remove the old version, then embed and store the new chunks
        vectorStore.delete(new FilterExpressionBuilder().eq("docId", docId).build());
        vectorStore.add(chunks);
        log.info("Ingested {} pages as {} chunks from {}", pages.size(), chunks.size(), docId);
    }
}

@Component
class IngestOnStartup implements ApplicationRunner {

    private final PolicyIngestionService ingestion;

    IngestOnStartup(PolicyIngestionService ingestion) {
        this.ingestion = ingestion;
    }

    @Override
    public void run(ApplicationArguments args) {
        ingestion.ingest(new ClassPathResource("docs/refund-policy.pdf"), "refund-policy");
        ingestion.ingest(new ClassPathResource("docs/student-handbook.pdf"), "student-handbook");
    }
}

@RestController
public class PolicyQaController {

    private final ChatClient chatClient;

    public PolicyQaController(ChatClient.Builder builder, VectorStore vectorStore) {
        this.chatClient = builder
            .defaultSystem("You answer questions for Webnest students using only the provided context. "
                + "If the context does not contain the answer, say you don't know.")
            .defaultAdvisors(QuestionAnswerAdvisor.builder(vectorStore)
                .searchRequest(SearchRequest.builder().topK(4).similarityThreshold(0.6).build())
                .build())
            .build();
    }

    @GetMapping("/policy/ask")
    public String ask(@RequestParam String q) {
        return chatClient.prompt().user(q).call().content();
    }
}

@Service
public class HandbookAssistant {

    private final ChatClient chatClient;

    public HandbookAssistant(ChatClient.Builder builder, VectorStore vectorStore) {
        Advisor rag = RetrievalAugmentationAdvisor.builder()
            .queryTransformers(RewriteQueryTransformer.builder()
                .chatClientBuilder(builder.build().mutate())
                .build())
            .documentRetriever(VectorStoreDocumentRetriever.builder()
                .vectorStore(vectorStore)
                .similarityThreshold(0.55)
                .topK(5)
                .build())
            .queryAugmenter(ContextualQueryAugmenter.builder()
                .allowEmptyContext(false)   // refuse instead of guessing when nothing relevant is found
                .build())
            .build();

        this.chatClient = builder.defaultAdvisors(rag).build();
    }

    public record Answer(String text, List<String> sources) {}

    public Answer ask(String question) {
        ChatResponse response = chatClient.prompt()
            .user(question)
            .call()
            .chatResponse();

        // The advisor stores the retrieved documents in the response metadata
        List<Document> docs = response.getMetadata().get(RetrievalAugmentationAdvisor.DOCUMENT_CONTEXT);
        List<String> sources = docs == null ? List.of() : docs.stream()
            .map(d -> d.getMetadata().get("docId") + " p." + d.getMetadata().get("page_number"))
            .distinct()
            .toList();

        return new Answer(response.getResult().getOutput().getText(), sources);
    }
}