Creating Personalized ChatGPT with Spring AI — Part 2

Personalized ChatGPT with Spring AI

In part one of this series, we explored the exciting potential of large language models (LLMs) and how Spring AI, a novel framework, empowers Java developers to harness this power for chatbot creation. We delved into the core concepts of LLMs and how Spring AI simplifies LLM integration. We also embarked on a step-by-step guide to coding a chatbot using Spring AI — but that journey isn’t over yet!

In this part we’ll cover the following:

• Refine our chatbot’s functionality by incorporating personalization techniques.
• Work with a new LLM called Mistral

Let’s Make Chatbots Smarter

Getting Smarter

In part one, we built a basic chatbot using the open-source LLM model Llama2. While it was a fun experiment, real-world chatbots need more! Imagine a chatbot that remembers your past conversations or tailors its responses to specific documents, like research papers or company reports.

This is exactly where personalization comes in. Buckle up, because in this part, we’ll transform our basic Spring application into a context-aware chatbot! We’ll explore how to make it answer questions specific to a PDF document. (For this example, we’ll use a research paper of Mistral — because, well, it’s a research paper about AI which will be used by an AI and I am a fan of inception!)

🚀 Without Further Ado: Let’s Dive In!

Prerequisites

• Java 17 or new versions of Java
• Ollama
• Basic knowledge of Spring
• Part 1 of this series

Let’s Code

1. Let’s add a new dependency which will allow us to read the pdf and convert that into document.

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-tika-document-reader</artifactId>
</dependency>

2. Download the research paper of Mistral as pdf ( feel free to add any other pdf if you want ) and copy it in resources folder of your spring application.

3. Create a new service for reading the file.

@Service
public class PdfFileReader {
    private final VectorStore vectorStore;

    @Value("classpath:researchPaper.pdf")
    public Resource pdfResource;

    public PdfFileReader(@Autowired VectorStore vectorStore) {
        this.vectorStore = vectorStore;
    }

    @PostConstruct
    public void init() {
        TikaDocumentReader tikaDocumentReader = new TikaDocumentReader(pdfResource);
        TokenTextSplitter tokenTextSplitter = new TokenTextSplitter();
        vectorStore.add(tokenTextSplitter.apply(tikaDocumentReader.get()));
    }
}

There are 2 important concepts which are introduced here. Let’s try to dissect them before we move forward.

1. TokenTextSplitter — If you have used any generative AI application then you might’ve heard about tokens. In the world of large language models (LLMs), tokens are the building blocks used to understand and generate text. They function like individual words, but can be differently handled in different models. That’s what we are also doing in this service, we are reading the file’s content and then tokenizing the content.
2. VectorStore — Before understanding this we should understand what are Vector databases. In Spring’s official documentation’s words ‘vector database is a specialized type of database that plays an essential role in AI applications’. In a nutshell, vector databases store data as points in multi-dimensional space, allowing for fast searches based on similarity instead of exact matches. Imagine you have a music streaming service and want to recommend songs to users. A traditional database might store information like genre or artist. But what if a user loves a particular song and wants to find similar ones? Here’s where a vector database comes in. It can store the music as vectors based on features like tempo, melody, and rhythm. When a user likes a song, the service can search the database for songs with similar vectors, recommending music that sounds alike, even if from a different genre or artist. Coming back to VectorStore, VectorStore is defines different ways to interact with Vector Databases. Spring AI has support for a lot of vector stores but for this example we are going to use SimpleVectorStore.

4. Update AIConfiguration.java to use with mistral model and create instance of VectorStore.

@Configuration
public class AIConfiguration {
    @Bean
    @Primary
    OllamaApi ollamaApi() {
        return new OllamaApi("http://localhost:11434/");
    }

    @Bean
    @Primary
    ChatClient chatClient(OllamaApi ollamaApi) {
        return new OllamaChatClient(ollamaApi).withModel("mistral")
                .withDefaultOptions(OllamaOptions.create()
                        .withModel(OllamaOptions.DEFAULT_MODEL)
                        .withTemperature(0.4f));
    }
    
    @Bean
    VectorStore vectorStore(EmbeddingClient embeddingClient) {
        return new SimpleVectorStore(embeddingClient);
    }
}

5. Update OllamaAiService.java and create a new method

public String searchInUserContext(String query) {
        List<Document> similarDocuments = vectorStore.similaritySearch(query);
        String information = similarDocuments.stream()
                .map(Document::getContent)
                .collect(Collectors.joining(System.lineSeparator()));
        SystemPromptTemplate systemPromptTemplate = new SystemPromptTemplate(
                """
                            You are a helpful assistant.
                            Use only the following information to answer the question.
                            Do not use any other information. If you do not know, simply answer: Unknown.

                            {information}
                        """);
        Message systemMessage = systemPromptTemplate
                                  .createMessage(Map.of("information", information));
        PromptTemplate userPromptTemplate = new PromptTemplate("{query}");
        Message userMessage = userPromptTemplate
                              .createMessage(Map.of("query", query));
        Prompt prompt = new Prompt(List
                            .of(systemMessage, userMessage));
        return chatClient.call(prompt)
                .getResult().getOutput().getContent();
    }

6. Create a new endpoint in LlamaRestController.java which invokes this new method and return the response.

@RestController
public class OllamaRestController {
    private final OllamaAiService llamaAiService;
    @Autowired
    public OllamaRestController(OllamaAiService llamaAiService) {
        this.llamaAiService = llamaAiService;
    }
    @GetMapping("chat")
    public ResponseEntity<String> generate(
            @RequestParam(value = "prompt",
             defaultValue = "Why don't you follow us?")
            String prompt) {
        String aiResponse = llamaAiService.generateMessage(prompt);
        return ResponseEntity.status(HttpStatus.OK).body(aiResponse);
    }
    @GetMapping("chat/joke/{topic}")
    public ResponseEntity<String> generateJoke(@RequestParam(value = "topic",
     defaultValue = "Tech")
            String topic) {
        String aiResponse = llamaAiService.generateJoke(topic);
        return ResponseEntity.status(HttpStatus.OK).body(aiResponse);
    }
    
    @GetMapping("search/{prompt}")
    public ResponseEntity<String> searchInUserContext(@RequestParam(value = "prompt",
             defaultValue = "Summarize the document for me")
            String prompt) {
        String aiResponse = llamaAiService.searchInUserContext(prompt);
        return ResponseEntity.status(HttpStatus.OK).body(aiResponse);
    }
}

7. Let’s run Ollama and bring up our LLM model. If you have installed Ollama then you can run it by the below command

ollama run mistral // i.e. ollama run <model>

If you are running Ollama as docker image then you can use the below command

docker run -d --gpus=all -v ollama:/root/.ollama -p 11434:11434 
--name ollama ollama/ollama

Once Ollama is running, test it out by running the below curl

curl --location '<http://localhost:11434/api/generate>' \\
--header 'Content-Type: application/json' \\
--data '{
  "model": "mistral",
  "prompt": "What is AI ?"
 }'

1. Now let’s bring up our Spring app and test the rest endpoints.

curl --location '<http://localhost:8080/search?prompt=Summarize the performance improvement and compare it with Llama2' \\
--data ''

Boom! Now your chatbot can answer questions based on the specific PDF you uploaded.

Congratulations! You’ve successfully transformed your basic Spring application into a context-aware chatbot that can answer questions based on a specific PDF document. This is a significant step towards building truly intelligent and engaging chatbots.

But the journey doesn’t end here! In the next part of this exciting series, we’ll venture into the fascinating realm of text-to-image generation. We’ll leverage Spring’s capabilities to explore how to create an application that can take textual descriptions and transform them into stunning visuals. Imagine a chatbot that can not only answer your questions about a document but also generate an infographic summarizing its key points — the possibilities are endless!

Follow us to not miss the next installment of the series.