Retrieval-Augmented Generation With MongoDB and Spring AI: Bringing AI to Your Java Applications

1
spring.application.name=RagApp  
2
  
3
spring.ai.openai.api-key=<Your-API-Key>
4
spring.ai.openai.chat.options.model=gpt-4o  
5
  
6
spring.ai.vectorstore.mongodb.initialize-schema=true  
7
  
8
spring.data.mongodb.uri=<Your-Connection-URI>
9
spring.data.mongodb.database=rag

1
import org.springframework.ai.embedding.EmbeddingModel;
2
import org.springframework.ai.openai.OpenAiEmbeddingModel;
3
import org.springframework.ai.openai.api.OpenAiApi;
4
import org.springframework.ai.vectorstore.MongoDBAtlasVectorStore;
5
import org.springframework.ai.vectorstore.VectorStore;
6
import org.springframework.beans.factory.annotation.Value;
7
import org.springframework.context.annotation.Bean;
8
import org.springframework.context.annotation.Configuration;
9
import org.springframework.data.mongodb.core.MongoTemplate;
10
11
@Configuration
12
public class Config {
13
14
    @Value("${spring.ai.openai.api-key}")
15
    private String openAiKey;
16
17
    @Bean
18
    public EmbeddingModel embeddingModel() {
19
        return new OpenAiEmbeddingModel(new OpenAiApi(openAiKey));
20
    }
21
22
    @Bean
23
    public VectorStore mongodbVectorStore(MongoTemplate mongoTemplate, EmbeddingModel embeddingModel) {
24
        return new MongoDBAtlasVectorStore(mongoTemplate, embeddingModel,
25
                MongoDBAtlasVectorStore.MongoDBVectorStoreConfig.builder().build(), true);
26
    }
27
28
}

1
package com.mongodb.RagApp.service;
2
3
import com.fasterxml.jackson.databind.ObjectMapper;
4
import org.springframework.ai.document.Document;
5
import org.springframework.ai.vectorstore.VectorStore;
6
import org.springframework.beans.factory.annotation.Autowired;
7
import org.springframework.core.io.ClassPathResource;
8
import org.springframework.stereotype.Service;
9
10
import java.io.BufferedReader;
11
import java.io.InputStream;
12
import java.io.InputStreamReader;
13
import java.util.ArrayList;
14
import java.util.List;
15
import java.util.Map;
16
17
@Service
18
public class DocsLoaderService {
19
20
    private static final int MAX_TOKENS_PER_CHUNK = 2000; 
21
    private final VectorStore vectorStore;
22
    private final ObjectMapper objectMapper;
23
24
    @Autowired
25
    public DocsLoaderService(VectorStore vectorStore, ObjectMapper objectMapper) {
26
        this.vectorStore = vectorStore;
27
        this.objectMapper = objectMapper;
28
    }
29
30
    public String loadDocs() {
31
        try (InputStream inputStream = new ClassPathResource("docs/devcenter-content-snapshot.2024-05-20.json").getInputStream();
32
             BufferedReader reader = new BufferedReader(new InputStreamReader(inputStream))) {
33
34
            List<Document> documents = new ArrayList<>();
35
            String line;
36
37
            while ((line = reader.readLine()) != null) {
38
                Map<String, Object> jsonDoc = objectMapper.readValue(line, Map.class);
39
                String content = (String) jsonDoc.get("body");
40
41
                // Split the content into smaller chunks if it exceeds the token limit
42
                List<String> chunks = splitIntoChunks(content, MAX_TOKENS_PER_CHUNK);
43
44
                // Create a Document for each chunk and add it to the list
45
                for (String chunk : chunks) {
46
                    Document document = createDocument(jsonDoc, chunk);
47
                    documents.add(document);
48
                }
49
                // Add documents in batches to avoid memory overload
50
                if (documents.size() >= 100) {
51
                    vectorStore.add(documents);
52
                    documents.clear();
53
                }
54
            }
55
            if (!documents.isEmpty()) {
56
                vectorStore.add(documents);
57
            }
58
59
            return "All documents added successfully!";
60
        } catch (Exception e) {
61
            return "An error occurred while adding documents: " + e.getMessage();
62
        }
63
    }
64
65
    private Document createDocument(Map<String, Object> jsonMap, String content) {
66
        Map<String, Object> metadata = (Map<String, Object>) jsonMap.get("metadata");
67
68
        metadata.putIfAbsent("sourceName", jsonMap.get("sourceName"));
69
        metadata.putIfAbsent("url", jsonMap.get("url"));
70
        metadata.putIfAbsent("action", jsonMap.get("action"));
71
        metadata.putIfAbsent("format", jsonMap.get("format"));
72
        metadata.putIfAbsent("updated", jsonMap.get("updated"));
73
74
        return new Document(content, metadata);
75
    }
76
77
    private List<String> splitIntoChunks(String content, int maxTokens) {
78
        List<String> chunks = new ArrayList<>();
79
        String[] words = content.split("\\s+");
80
        StringBuilder chunk = new StringBuilder();
81
        int tokenCount = 0;
82
83
        for (String word : words) {
84
            // Estimate token count for the word (approximated by character length for simplicity)
85
            int wordTokens = word.length() / 4;  // Rough estimate: 1 token = ~4 characters
86
            if (tokenCount + wordTokens > maxTokens) {
87
                chunks.add(chunk.toString());
88
                chunk.setLength(0); // Clear the buffer
89
                tokenCount = 0;
90
            }
91
            chunk.append(word).append(" ");
92
            tokenCount += wordTokens;
93
        }
94
        if (chunk.length() > 0) {
95
            chunks.add(chunk.toString());
96
        }
97
        return chunks;
98
    }
99
}

1
import com.mongodb.RagApp.service.DocsLoaderService;
2
import org.springframework.web.bind.annotation.GetMapping;
3
import org.springframework.web.bind.annotation.RequestMapping;
4
import org.springframework.web.bind.annotation.RestController;
5
6
@RestController
7
@RequestMapping("/api/docs")
8
public class DocsLoaderController {
9
10
    private DocsLoaderService docsLoaderService;
11
12
    public DocsLoaderController(DocsLoaderService docsLoaderService) {
13
        this.docsLoaderService = docsLoaderService;
14
    }
15
16
    @GetMapping("/load")
17
    public String loadDocuments() {
18
        return docsLoaderService.loadDocs();
19
    }
20
21
}

1
import org.springframework.ai.chat.client.ChatClient;  
2
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;  
3
import org.springframework.ai.vectorstore.SearchRequest;  
4
import org.springframework.ai.vectorstore.VectorStore;  
5
import org.springframework.web.bind.annotation.CrossOrigin;  
6
import org.springframework.web.bind.annotation.GetMapping;  
7
import org.springframework.web.bind.annotation.RequestParam;  
8
import org.springframework.web.bind.annotation.RestController;
9
10
@RestController
11
public class RagController {
12
13
    private final ChatClient chatClient;
14
15
    public RagController(ChatClient.Builder builder, VectorStore vectorStore) {
16
        this.chatClient = builder
17
                .defaultAdvisors(new QuestionAnswerAdvisor(vectorStore, SearchRequest.defaults()))
18
                .build();
19
    }
20
21
    @GetMapping("/question")
22
    public String question(@RequestParam(value = "message", defaultValue = "How to analyze time-series data with Python and MongoDB?") String message) {
23
        return chatClient.prompt()
24
                .user(message)          
25
                .call()                 
26
                .content();
27
    }
28
}

1
private static final String DEFAULT_USER_TEXT_ADVISE = """
2
			Context information is below.
3
			---------------------
4
			{question_answer_context}
5
			---------------------
6
			Given the context and provided history information and not prior knowledge,
7
			reply to the user comment. If the answer is not in the context, inform
8
			the user that you can't answer the question.
9
			""";