Build a RAG Chatbot with Spring Boot

@RestController
@RequestMapping("/api/documents")
public class DocumentController {
 
    private final DocumentIngestionService ingestionService;
 
    public DocumentController(DocumentIngestionService ingestionService) {
        this.ingestionService = ingestionService;
    }
 
    @PostMapping("/upload")
    public ResponseEntity<Map<String, Object>> uploadDocument(
            @RequestParam("file") MultipartFile file,
            @RequestParam(value = "category", defaultValue = "general") String category) {
 
        if (file.isEmpty()) {
            return ResponseEntity.badRequest()
                .body(Map.of("error", "No file provided"));
        }
 
        String filename = file.getOriginalFilename();
        if (filename == null || !filename.endsWith(".pdf")) {
            return ResponseEntity.badRequest()
                .body(Map.of("error", "Only PDF files are supported"));
        }
 
        try {
            int chunkCount = ingestionService.ingestPdf(file, category);
            return ResponseEntity.ok(Map.of(
                "message", "Document indexed successfully",
                "file", filename,
                "chunks_created", chunkCount
            ));
        } catch (IOException e) {
            return ResponseEntity.internalServerError()
                .body(Map.of("error", "Failed to process document: " + e.getMessage()));
        }
    }
}

@Service
public class RagChatService {
 
    private final VectorStore vectorStore;
    private final ChatClient chatClient;
    private final ChatSessionRepository sessionRepo;
    private final ChatMessageRepository messageRepo;
 
    private static final String SYSTEM_PROMPT = """
        You are a helpful assistant that answers questions based on
        the provided context documents.
 
        Rules:
        - Answer ONLY using the information in the context below.
        - If the context does not contain enough information to answer
          the question, say: "I could not find that information in the
          available documents. Please try rephrasing your question."
        - Do not make up facts or use information from your training data.
        - If the user asks a follow-up question, use the chat history
          to understand what they are referring to.
        - Be concise and direct. Avoid unnecessary filler.
        """;
 
    public ChatResponse chat(UUID sessionId, String userQuestion) {
 
        // Step 1: Get or create the session
        ChatSession session = sessionRepo.findById(sessionId)
            .orElseGet(() -> sessionRepo.save(new ChatSession()));
 
        // Step 2: Retrieve relevant chunks from pgvector
        List<Document> relevantDocs = vectorStore.similaritySearch(
            SearchRequest.query(userQuestion)
                .withTopK(4)
                .withSimilarityThreshold(0.65)
        );
 
        // Step 3: Build the context string
        String context = relevantDocs.stream()
            .map(doc -> {
                String source = (String) doc.getMetadata()
                    .getOrDefault("source", "Unknown");
                return "[Source: " + source + "]\n" + doc.getContent();
            })
            .collect(Collectors.joining("\n\n---\n\n"));
 
        if (context.isBlank()) {
            context = "No relevant documents were found for this query.";
        }
 
        // Step 4: Build conversation history
        List<ChatMessage> history = messageRepo
            .findBySessionOrderByCreatedAtAsc(session);
 
        // Step 5: Assemble and call the LLM
        ChatClient.ChatClientRequest request = chatClient.prompt()
            .system(SYSTEM_PROMPT)
            .user(u -> u.text(
                "Context:\n{context}\n\nQuestion: {question}")
                .param("context", context)
                .param("question", userQuestion)
            );
 
        // Inject conversation history
        for (ChatMessage msg : history) {
            if ("user".equals(msg.getRole())) {
                request = request.messages(
                    new UserMessage(msg.getContent()));
            } else {
                request = request.messages(
                    new AssistantMessage(msg.getContent()));
            }
        }
 
        String answer = request.call().content();
 
        // Step 6: Persist the exchange
        saveMessage(session, "user", userQuestion);
        saveMessage(session, "assistant", answer);
 
        // Step 7: Return response with sources
        List<String> sources = relevantDocs.stream()
            .map(d -> (String) d.getMetadata()
                .getOrDefault("source", "Unknown"))
            .distinct()
            .collect(Collectors.toList());
 
        return new ChatResponse(session.getId(), answer, sources);
    }
 
    private void saveMessage(ChatSession session, String role, String content) {
        ChatMessage msg = new ChatMessage();
        msg.setSession(session);
        msg.setRole(role);
        msg.setContent(content);
        messageRepo.save(msg);
    }
}