How to Manage Multi-User Chat Memory with AI Services using ChatMemoryProvider?

Integrating AI into applications involves handling context-aware conversations. When working with multi-user systems, managing separate conversational memory per user is crucial. This post explores how to use ChatMemoryProvider with AiServices to maintain individual user sessions, prevent memory leaks, and ensure thread safety in chat-based AI services.

Problem Statement

When building conversational applications, maintaining conversation history (memory) is essential for contextual responses. By default, a shared ChatMemory instance can result in cross-user context conflicts not acceptable in multi-user environments.

 

Solution: Use ChatMemoryProvider

Instead of sharing a single ChatMemory, use the @MemoryId annotation to provide user-specific memory.

 

AI Service Interface

public interface ChatAssistant {
  String chat(@MemoryId int memoryId, @UserMessage String message);
}

Service Instantiation with Memory Provider

chatAssistant = AiServices.builder(ChatAssistant.class)
    .chatModel(chatModel)
    .chatMemoryProvider(memoryId -> MessageWindowChatMemory.withMaxMessages(10))
    .build();

Handling Multiple Users

String answerToRam = chatAssistant.chat(1, "Hello, my name is Ram");
String answerToKrishna = chatAssistant.chat(2, "Hello, my name is Krishna");

Each user gets an isolated memory instance identified by their unique memoryId.

 

When using ChatMemory in this way, it's important to evict memory associated with conversations that are no longer needed to prevent memory leaks. To access and manage these internal chat memories, simply ensure that the interface defining your AI service extends ChatMemoryAccess.

public interface ChatAssistant extends ChatMemoryAccess{
  String chat(@MemoryId int memoryId, @UserMessage String message);
}

This makes it possible to both access the ChatMemory instance of a single conversation and to get rid of it when the conversation is terminated.  

 

List<ChatMessage> messagesWithKlaus = chatAssistant.getChatMemory(1).messages();
boolean chatMemoryWithFrancineEvicted = chatAssistant.evictChatMemory(2);

Find the below working application.

 

Step 1: Create new maven project chatmemory-demo

 

Step 2: Update pom.xml with maven dependencies.

 

pom.xml

 

<project xmlns="http://maven.apache.org/POM/4.0.0"
  xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
  xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 https://maven.apache.org/xsd/maven-4.0.0.xsd">
  <modelVersion>4.0.0</modelVersion>
  <groupId>com.sample.app</groupId>
  <artifactId>chatmemory-demo</artifactId>
  <version>0.0.1-SNAPSHOT</version>
  
  <properties>
    <maven.compiler.source>21</maven.compiler.source>
    <maven.compiler.target>21</maven.compiler.target>
    <java.version>21</java.version>
  </properties>

  <parent>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-parent</artifactId>
    <version>3.3.10</version>
  </parent>

  <dependencyManagement>
    <dependencies>
      <dependency>
        <groupId>dev.langchain4j</groupId>
        <artifactId>langchain4j-bom</artifactId>
        <version>1.0.1</version>
        <type>pom</type>
        <scope>import</scope>
      </dependency>
    </dependencies>
  </dependencyManagement>

  <dependencies>
    <dependency>
      <groupId>dev.langchain4j</groupId>
      <artifactId>langchain4j-spring-boot-starter</artifactId>
    </dependency>


    <dependency>
      <groupId>dev.langchain4j</groupId>
      <artifactId>langchain4j-ollama-spring-boot-starter</artifactId>
    </dependency>

    <dependency>
      <groupId>org.springframework.boot</groupId>
      <artifactId>spring-boot-starter-web</artifactId>
    </dependency>

    <dependency>
      <groupId>org.springdoc</groupId>
      <artifactId>springdoc-openapi-starter-webmvc-ui</artifactId>
      <version>2.6.0</version>
    </dependency>

    <!--
    https://mvnrepository.com/artifact/jakarta.validation/jakarta.validation-api -->
    <dependency>
      <groupId>jakarta.validation</groupId>
      <artifactId>jakarta.validation-api</artifactId>
    </dependency>
  </dependencies>

  <build>
    <plugins>
      <plugin>
        <groupId>org.springframework.boot</groupId>
        <artifactId>spring-boot-maven-plugin</artifactId>
        <executions>
          <execution>
            <goals>
              <goal>repackage</goal> <!-- Important -->
            </goals>
          </execution>
        </executions>
      </plugin>
    </plugins>
  </build>

</project>

Step 3: Define ChatResponse dto.

 

ChatResponse.java

 

package com.sample.app.dto;

public class ChatResponse {
  private String message;

  public ChatResponse() {
  }

  public ChatResponse(String message) {
    this.message = message;
  }

  public String getMessage() {
    return message;
  }

  public void setMessage(String message) {
    this.message = message;
  }
}

 

Step 4: Define ChatAssistant.

 

ChatAssistant.java

 

package com.sample.app.interfaces;

import dev.langchain4j.service.MemoryId;
import dev.langchain4j.service.UserMessage;
import dev.langchain4j.service.memory.ChatMemoryAccess;

public interface ChatAssistant extends ChatMemoryAccess{
  String chat(@MemoryId int memoryId, @UserMessage String message);
}

Step 5: Define SwaggerConfig and LangchainConfig classes.

 

SwaggerConfig.java

 

package com.sample.app.config;

import org.springframework.context.annotation.Configuration;

import io.swagger.v3.oas.annotations.OpenAPIDefinition;
import io.swagger.v3.oas.annotations.info.Info;

@Configuration
@OpenAPIDefinition(info = @Info(title = "Chat service Application", version = "v1"))
public class SwaggerConfig {

}

 

LangchainConfig.java

package com.sample.app.config;

import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import dev.langchain4j.http.client.spring.restclient.SpringRestClientBuilderFactory;
import dev.langchain4j.model.ollama.OllamaChatModel;

@Configuration
public class LangchainConfig {

  @Bean
  public OllamaChatModel ollamaLanguageModel() {
    return OllamaChatModel.builder().baseUrl("http://localhost:11434")
        .modelName("llama3.2")
        .logRequests(true)
        .httpClientBuilder(new SpringRestClientBuilderFactory().create()) // explicitly use Spring's HTTP client
        .build();
  }
}

Step 6: Define ChatService.

 

ChatService.java

 

package com.sample.app.service;

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

import com.sample.app.dto.ChatResponse;
import com.sample.app.interfaces.ChatAssistant;

import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.ollama.OllamaChatModel;
import dev.langchain4j.service.AiServices;

@Service
public class ChatService {
  @Autowired
  private OllamaChatModel chatModel;

  private ChatAssistant chatAssistant;

  public ChatResponse chat(Integer memoryId, String userMessage) {
    if (chatAssistant == null) {
      synchronized (ChatService.class) {
        if (chatAssistant == null) {
          chatAssistant = AiServices.builder(ChatAssistant.class).chatModel(chatModel)
              .chatMemoryProvider(chatMemoryId -> MessageWindowChatMemory.withMaxMessages(10)).build();
        }

      }

    }

    String story = chatAssistant.chat(memoryId, userMessage);
    return new ChatResponse(story);
  }
}

 

Step 7: Define ChatController class.

 

ChatController.java

 

package com.sample.app.controller;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.CrossOrigin;
import org.springframework.web.bind.annotation.PostMapping;
import org.springframework.web.bind.annotation.RequestBody;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import com.sample.app.dto.ChatResponse;
import com.sample.app.service.ChatService;

import io.swagger.v3.oas.annotations.tags.Tag;
import jakarta.validation.Valid;
import jakarta.validation.constraints.NotEmpty;
import jakarta.validation.constraints.Positive;

@RestController
@RequestMapping("/api/chat")
@CrossOrigin("*")
@Tag(name = "Chat Controller", description = "This section contains APIs related to Chat APIs Powered by Ollama")
public class ChatController {

  @Autowired
  private ChatService chatService;

  @PostMapping
  public ChatResponse chat(@RequestBody @Valid ChatRequestBody chatRequestBody) {
    return chatService.chat(chatRequestBody.chatId, chatRequestBody.getMessage());
  }

  private static class ChatRequestBody {
    @NotEmpty
    private String message;

    @Positive
    private Integer chatId;

    public String getMessage() {
      return message;
    }

    public void setMessage(String message) {
      this.message = message;
    }

    public Integer getChatId() {
      return chatId;
    }

    public void setChatId(Integer chatId) {
      this.chatId = chatId;
    }

  }
}

Step 8: Define main application class.

 

App.java

package com.sample.app;

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

@SpringBootApplication
public class App {

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

}

 

Build the project

Navigate to the project root directory where pom.xml is located and execute following command to generate the artifact.

 

mvn clean install

Upon successful execution of the command, you can see chatmemory-demo-0.0.1-SNAPSHOT.jar file in the target folder.

$ ls ./target/
chatmemory-demo-0.0.1-SNAPSHOT.jar    generated-test-sources
chatmemory-demo-0.0.1-SNAPSHOT.jar.original maven-archiver
classes           maven-status
generated-sources       test-classes

 

Run the Application

Execute following command to run the Application.

 

java -jar ./target/chatmemory-demo-0.0.1-SNAPSHOT.jar --server.port=1235

 

Open the below url in browser to interact with swagger endpoint.

http://localhost:1235/swagger-ui/index.html

 

Test the Application.

 

Step 1: Run the API /api/chat with below payload.

{
  "message": "Hi My name is Ram",
  "chatId": 1
}

Step 2: Run the API /api/chat with below payload.

{
  "message": "Hi My name is Krishna",
  "chatId": 2
}

Step 3: Let’s ask the LLM, what is my name for the chatId 2.

 

Execute the API /api/chat with below payload.

 

{
  "message": "Guess My Name?",
  "chatId": 2
}

I received the response like below.

{
  "message": "That sounds like a fun game!\n\nTo start, I'll ask some yes or no questions to try and guess your name. Here's my first question:\n\nIs your name Krishna?"
}

 

You can download the application from this link.  

Previous                                                    Next                                                    Home