AI Engineering in Java | Part 3

AI Engineering in Java | Part 3

Previous titles in this series:

In our third installation of this ongoing series I only planned to look at incorporating some chat memory functions to store the chats for the users. But I always like to have my projects – although they're “just” concepts or playgrounds – to be as realistic as it gets. So what is any storing of user chats useful without any users? And for what are any users useful if they can't log in and get proper authorization, right? As for this Quarkus makes it really easy for us to set this up with Keycloak as identity management provider. So here's the plan for this episode:

  1. Add basic authorization and authentication with Keycloak dev services and Quarkus's role based access capabilities.
  2. Connect the Postgresql database and create a simple repository to store and retrieve user chats using plain old JDBC.
  3. Expand our UI, to get access to former chats and create new chats for logged in users.

Still there's a little disclaimer to be announced: This part is most likely the one with the least AI specific content, as it's more about wiring up some components and creating some user experience in our UI. Also: some components will drop in later parts, when we incorporate greater abstractions and harness the interfaces which Langchain4J gives us.


Keycloak as OIDC provider

First let's add a Keycloak server as OIDC provider and for our user- and identity-management system. Quarkus has built-in capabilities to easily create a Keycloak server at startup, creating realms and clients and also adding two users: alice (who gets admin and user roles) and bob (who gets just a user) role. A minimalistic setup but perfect for experimentation and constructing proof of concepts. The best part? Just add a few lines to /src/main/resources/application.properties as we will do now:

quarkus.keycloak.devservices.enabled=true
quarkus.keycloak.devservices.realm-name=javaai
quarkus.keycloak.devservices.create-realm=true
quarkus.keycloak.devservices.create-client=true
quarkus.oidc.client.id=quarkus
quarkus.oidc.credentials.secret=mytopsecret

After that add the appropriate Quarkus oidc extension as a dependency to your pom.xml:

<dependencies>
...
        <dependency>
            <groupId>io.quarkus</groupId>
            <artifactId>quarkus-oidc</artifactId>
        </dependency>
...
</dependencies>

When you startup Quarkus in dev-mode now via mvn quarkus:dev, Quarkus is pulling the latest Keycloak docker, installs it, sets up a realm called “javaai” and installs a client “quarkus” with the secret “mytopsecret”. Of course this is all just for developing purposes! 😉 But it works very smoothly.

As you can see: the Keycloak OIDC provider was added to our dev-ui.

Keycloak is a great but also complex tool for OIDC, I use regularly and sucessfully in my projects. So I can't get into any details in this article. But if you are curious, I recommend starting with some video-tutorials by Niko Köbler, who is called “Mr. Keycloak” for reasons!

The cool thing is, Quarkus set up everything important for us to just start using Keycloaks capabilities. Just klick on “Keycloak Admin” in your dev-ui and login with “admin” for user and “admin” for password, too.

For the protected routes to work in Quarkus, we have to define them in the application.properties, too.

quarkus.oidc.application-type=web-app

# role access for normal users to authenticate
quarkus.http.auth.permission.authenticated.paths=/user*
quarkus.http.auth.permission.authenticated.policy=authenticated

# role access for admins
quarkus.http.auth.permission.admin.paths=/admin*
quarkus.http.auth.permission.admin.policy=admin-role-policy
quarkus.http.auth.policy.admin-role-policy.roles-allowed=admin


quarkus.oidc.logout.path=/logout
quarkus.oidc.logout.post-logout-path=/

With the application-type=web-app we tell Quarkus to enable authentication for routes. Then we add the protected routes. Everything under the route /user is protected and users have to authenticate to view the pages. Under /admin and beyond, only users with the role admin are allowed to view the pages. We also defined a default logout path. Let's try to add two views to test the admin and the authenticated view:

// src/main/java/de/datenschauer/ui/UserView.java

@Route("user")
public class UserView extends VerticalLayout {

  public UserView() {
    setSizeFull();
    add(new Paragraph("Only logged in users can see this!"));
  }

}
// src/main/java/de/datenschauer/ui/AdminView.java

@Route("admin")
public class AdminView extends VerticalLayout {

  public AdminView() {
    setSizeFull();
    add(new Paragraph("Only admins are allowed to see this!"));
  }
}

Restart Quarkus and then head to http://localhost:8080/user. You should be redirected to the Keycloak login page:

And when you're logged in you should see the user's page:

After that, head to http://localhost:8080/admin and you should get a HTTP 403 “unauthorized” error. First logout via http://localhost:8080/logout and then again try to navigate to the admin test-page. You should be able to log in again now as “alice” with password “alice”; and since “alice” has the role admin defined in Keycloak you should see the admin page.

That's nice and easy, isn't it? In the background Quarkus is storing session cookies, which we can also access in our code, to get some identity values for the logged in users. This will help us to associate chats with the users and store them accordingly later. Quarkus uses JWT (Json Web Token) for this. Here is, how we can access these values in a secure manner:

public class AdminView extends VerticalLayout {

  @Inject
  @IdToken
  JsonWebToken idToken;

...

  @PostConstruct
  private void init() {

    if (idToken != null) {
      var userName = idToken.getClaim("upn");
      var userId = idToken.getClaim("sub");

      add(new H2("User Info"));
      add(new Paragraph("Your are logged in as user >>" + userName + "<< with id: " + userId));
    }
  }
}

We injected the identity token provided by Keycloak. As there are two different such tokens, we have to tell Quarkus that we want the IdToken not the AccessToken via the @IdToken annotation. We also have to provide a @PostConstruct annotation for Vaadin; because when the component is rendered the idToken is still not available. Every token issued by an OIDC server like Keycloak has standard claims like "upn" for username or "sub" for internal ID. You have to get used to these abbreviations, but you could have a look at all the provided claims via idToken.getClaimNames().forEach(System.out::println); or so. But nevertheless, idToken helps us to build a chat storage next.


Add storage

I chose to talk to my database via plain old JDBC to not have to go into deeper subjects like ORMs, creating entities and so on (although I am quite sure most readers are familiar with those topics). For all the code to come, I will only post snippets as there's a lot going on. So to follow, you should really check out the repository at Codeberg!

Postgres Devservices

Quarkus has some really nice devservices which make it easy to set up a default PostgreSQL Database like we did just before with our Keycloak service. Just add the following to your application.properties file:

quarkus.datasource.db-kind=postgresql
quarkus.datasource.devservices.db-name=postgres
quarkus.datasource.devservices.port=52320
quarkus.datasource.devservices.username=quarkus
quarkus.datasource.devservices.password=quarkus
# initialization script for quarkus devservices
quarkus.datasource.devservices.init-script-path=init.sql
# Force container recreation when schema changes
quarkus.datasource.devservices.reuse=false

As you can see, we use a small init.sql file which creates our tables when the devservices start up. The content looks like the following:

CREATE TABLE IF NOT EXISTS userchats (
    chat_id SERIAL PRIMARY KEY ,
    user_id UUID NOT NULL ,
    created_at TIMESTAMPTZ NOT NULL
);

CREATE TABLE IF NOT EXISTS chatmessages (
    message_id SERIAL PRIMARY KEY ,
    chat_id INTEGER NOT NULL ,
    message TEXT NOT NULL ,
    type TEXT NOT NULL ,
    created_at TIMESTAMPTZ NOT NULL ,
    FOREIGN KEY (chat_id) REFERENCES userchats(chat_id) ,
    CONSTRAINT valid_msg_type CHECK ( type IN ('SYSTEM', 'USER', 'AI', 'TOOL_EXECUTION_RESULT', 'CUSTOM') )
);

Put the file inside your /src/main/resources folder. When Quarkus starts up, it creates two tables: one for holding the individual chats per user and one for the individual messages per userchat. We also do a check on message type, since we will only allow the standard message types allowed by Langchain4J and so we can later map them back to our ChatMessage objects.

Interactions

I created a little ChatMemoryService inside /src/main/java/de/datenschauer/db to do some basic CRUD operations:

public interface ChatMemoryService {
  int saveUserChat(UserChatDTO userChat) throws SQLException;

  void saveChatMessage(ChatMessageDTO message) throws SQLException;

  void deleteChatMessagesById(int chatId) throws SQLException;

  List<UserChatDTO> getChatsForUser(UUID userId) throws SQLException;

  List<ChatMessageDTO> getMessagesForChat(Integer chatId) throws SQLException;
}

You can find the Postgres implementation in /src/main/java/de/datenschauer/db/PostgresChatMemoryImpl.java. Nothing really special there. Just to mention, that again, Quarkus makes it easy for us to inject access to our PostgreSQL datasource:

@ApplicationScoped
public class PostgresChatMemoryImpl implements ChatMemoryService {

  @Inject
  DataSource dataSource;

  @Override
  public int saveUserChat(final UserChatDTO userChat) throws SQLException {
    var q = ...

    try(var conn = this.dataSource.getConnection(); var stmt = conn.prepareStatement(q)) {

      ...

Langchain4J MemoryStore

Langchain4J provides a simple interface to implement a MemoryStore which interacts with the underlying storage (Postgres in our case):

public interface ChatMemoryStore {


    List<ChatMessage> getMessages(Object memoryId);


    void updateMessages(Object memoryId, List<ChatMessage> messages);


    void deleteMessages(Object memoryId);
}

So we have to tell our implementation, how to retrieve message, how to update message and optionally how to delete messages. Keep in mind, that there are two concepts when we talk about memory. There is the concept of history and than there's just memory, as the Langchain4J docs explain:

  • History keeps all messages between the user and AI intact. History is what the user sees in the UI. It represents what was actually said.
  • Memory keeps some information, which is presented to the LLM to make it behave as if it "remembers" the conversation. Memory is quite different from history. Depending on the memory algorithm used, it can modify history in various ways: evict some messages, summarize multiple messages, summarize separate messages, remove unimportant details from messages, inject extra information (e.g., for RAG) or instructions (e.g., for structured outputs) into messages, and so on.

You can find the Postgres implementation of the ChatMemoryStore in /src/main/java/de/datenschauer/db/PostgresChatMemoryStore.java. We use two simple java records as our DTOs:

public record UserChatDTO(
    Integer chatId,
    UUID userId,
    OffsetDateTime createdAt
) {
}

public record ChatMessageDTO(
    Integer messageId,
    Integer chatId,
    String message,
    ChatMessageType type,
    OffsetDateTime createdAt
) {
}

With these we can map from database entries to and from ChatMessages correctly:

private ChatMessage createChatMessageFromDTO(final ChatMessageDTO dto) {
    return switch (dto.type()) {
      case SYSTEM ->  SystemMessage.from(dto.message());
      case USER -> UserMessage.from(dto.message());
      default -> AiMessage.from(dto.message());
    };
  }

  public int createNewUserChat(UUID userId) throws SQLException {
    return memory.saveUserChat(new UserChatDTO(null, userId, OffsetDateTime.now()));
  }

  @Override
  public List<ChatMessage> getMessages(Object chatId) {
    try {
      return memory.getMessagesForChat((int) chatId)
          .stream()
          .map(this::createChatMessageFromDTO)
          .toList();
    } catch (SQLException e) {
      log.error(e.getMessage());
    }
    return Collections.emptyList();
  }

  @Override
  public void updateMessages(Object chatId, List<ChatMessage> messages) {
    messages.forEach(message -> {
      String msgText = switch (message) {
        case UserMessage u -> u.hasSingleText() ? u.singleText() : u.contents().stream()
            .map(Object::toString).collect(Collectors.joining(" "));
        case AiMessage a -> a.text();
        default -> "";
      };
      final ChatMessageDTO messageDTO = new ChatMessageDTO(
          null, (int) chatId, msgText, message.type(), OffsetDateTime.now()
      );
      try {
        memory.saveChatMessage(messageDTO);
      } catch (SQLException e) {
        log.error(e.getMessage());
      }
    });
  }

There is also a ChatMessageSerializer but for learning purposes I found it easier and cleaner to work with DTOs and to get to know the underlying concepts of the ChatMessages implementations, like that a UserMessage hat contents, because not only text can be transferred to the LLM, but also audio and visual content for multimodal LLMs.


UI restructuring

Since we now have the possibility to save and restore user chats, we have to adopt our Vaadin UI to:

  • Get the UserID from the token of the logged in user.
  • Create a new chat for the user, whenever the page is entered the first time and the user issues a chat message, or when the “new” button is clicked.
  • Load and show the past userchats in a sidebar.
  • Load an individual chat when it is clicked.
  • Update a total chat history together with a chat memory for the current chat.

The new UI layout will look like this:

This is all happening now in the /src/main/java/de/datenschauer/ui/chat/ChatLayout.java class. I also added a ChatService.java class for some convenience functions of loading user chats and recreating the messages in the message view.

To get the correct user token and id, we also need to add the /chat url to our policies in application.properties so to always log in when using the chat functionality:

quarkus.http.auth.permission.authenticated.paths=/user*,/chat*
quarkus.http.auth.permission.authenticated.policy=authenticated

Memory and History

The most notable things besides the Vaadin components are the chatHistory and the chatMemory, where chatHistory stores all messages but chatMemory is an implementation of Langchain4J's MessageWindowChatMemory. It also interacts with the database with the PostgresChatMemoryStore but looks that the context, the LLM can hold doesn't get too big. So it evicts certain messages when the chatHistory is too large and only hands the latest messages over to the LLM to create something like a memory.

private ChatMemory buildChatMemory(int chatId) {
    return MessageWindowChatMemory.builder()
        .chatMemoryStore(store)
        .id(chatId)
        .maxMessages(ChatLayout.MAX_MESSAGES_IN_MEMORY)
        .build();
  }

When we interact with our model now, we can see, that the user chats are correctly stored in the database via Quarkus's Dev-UI:


Next up...

In our next installation of this series we will look into more acstractions and how to incorporate the Langchain4j AI Services.