Agente de IA confiável em prod com Java + Quarkus + Langchain4j - Parte 2 - Memória

Autores

@herbertbeckman - LinkedIn
@rndtavares - LinkedIn

Partes do artigo

1. Agente de IA confiável em prod com Java + Quarkus + Langchain4j - Parte 1 - AI as Service

2. Agente de IA confiável em prod com Java + Quarkus + Langchain4j - Parte 2 - Memória (este artigo)

3. Agente de IA confiável em prod com Java + Quarkus + Langchain4j - Parte 3 - RAG (em breve)

4. Agente de IA confiável em prod com Java + Quarkus + Langchain4j - Parte 4 - Guardrails (em breve)

Introdução

Quando criamos um agente, devemos ter em mente que os LLMs não guardam nenhum tipo de informação, ou seja, são stateless. Para que o nosso agente tenha a capacidade de se "lembrar" das informações, devemos implementar o gerenciamento de memória. O Quarkus já nos entrega uma memória padrão configurada, o porém disso é que ela pode literalmente derrubar o seu agente estourando a memória ram disponibilizada pra ele, como descrito nesta documentação do Quarkus, caso não se tome os devidos cuidados. Pra não termos mais esse problema e também para que seja possível utilizarmos nosso agente em um ambiente de escalabilidade, precisamos de um ChatMemoryStore.

Conceitos

Utilizamos um chat para interagir com o nosso agente e há conceitos importantes que devemos conhecer para que a nossa interação com ele possa ocorrer da melhor forma possível e não ocasione bugs em produção. Primeiramente precisamos conhecer os tipos de mensagens que utilizamos na hora de interagir com ele, são eles:

• Mensagens do usuário (UserMessage): A mensagem ou solicitação enviada pelo cliente final. Quando enviamos a mensagem no DevUI do quarkus, sempre estamos enviando uma UserMessage. Além disso, ela também é utilizada nos resultados de chamadas das ferramentas (tools) que vimos antes.

• Mensagens da IA (AiMessage): A mensagem de resposta do modelo. Sempre que o LLM responder pro nosso agente, ele receberá uma mensagem desse tipo. Este tipo de mensagem fica alternando o seu conteúdo entre uma resposta textual e solicitações de execução de ferramentas (tools).

• Mensagem do Sistema (SystemMessage): Esta mensagem pode ser definida somente 1 vez e é somente em tempo de desenvolvimento.

Agora que você conhece os 3 tipos de mensagens que temos, vamos explicar como elas devem se comportar com alguns gráficos. Todos os gráficos foram tirados da apresentação Java meets AI: Build LLM-Powered Apps with LangChain4j by Deandrea, Andrianakis, Escoffier, recomendo demais o vídeo.

O primeiro gráfico demonstra o uso dos 3 tipos de mensagens. UserMessage em azul, SystemMessage em vermelho e AiMessage em verde.

Neste segundo gráfico, demonstra-se como que a "memória" deve ser gerenciada. Um detalhe interessante é que devemos manter uma certa ordem nas mensagens e algumas premissas devem ser respeitadas.

• Só deve existir 1 mensagem do tipo SystemMessage;
• Após a SystemMessage, as mensagens sempre devem alternar entre UserMessage e AiMessage, nesta ordem. Se tivermos uma AiMessage após outra AiMessage, tomaremos uma exceção. O mesmo vale pra UserMessage seguidas.

Outro detalhe importante que você deve se atentar é sobre o tamanho do seu ChatMemory. Quanto maior a memória da sua interação, maior os custos com tokens, pois o LLM precisará processar mais texto pra dar uma resposta. Então estabeleça uma janela de memória que melhor se adequar pro seu caso de uso. Uma dica é verificar a média de mensagens dos seus clientes para ter uma ideia de tamanho de interação. Iremos mostrar a implementação através da MessageWindowChatMemory, a classe especializada em gerenciar isso pra gente no Langchain4j.

Agora que conhecemos todos esses conceitos e premissas, vamos por a mão na massa!

Configurando nosso ChatMemoryStore

Aqui vamos utilizar o MongoDB como um ChatMemoryStore. Utilizamos a doc do MongoDB e subimos uma instância no docker. Sinta-se a vontade pra configurar ele como bem desejar.

Adicionando nossa conexão com o MongoDB

Vamos iniciar adicionando a dependência necessária para termos uma conexão com o MongoDB utilizando o Quarkus.

<dependency>
    <groupId>io.quarkus</groupId>
    <artifactId>quarkus-mongodb-panache</artifactId>
</dependency>

Após as dependências, precisamos adicionar as configurações de conexão no nosso src/main/resources/application.properties.

quarkus.mongodb.connection-string=mongodb://${MONGODB_USER}:${MONGODB_PASSWORD}@localhost:27017
quarkus.mongodb.database=chat_memory

Ainda não iremos conseguir testar a nossa conexão com a base, pois antes precisamos criar nossas entidades e repositórios.

Criando nossa entidade e nosso repositório

Agora vamos implementar nossa entidade Interaction. Essa entidade terá a nossa lista de mensagens realizadas. Sempre que um cliente novo se conectar, será gerada uma nova Interaction. Se precisarmos reaproveitar essa Interaction, basta informamos o mesmo identificador da Interaction.

package <seupacote>;

import dev.langchain4j.data.message.ChatMessage;
import io.quarkus.mongodb.panache.common.MongoEntity;
import org.bson.codecs.pojo.annotations.BsonId;

import java.util.List;
import java.util.Objects;

@MongoEntity(collection = "interactions")
public class InteractionEntity {

    @BsonId
    private String interactionId;
    private List<ChatMessage> messages;

    public InteractionEntity() {
    }

    public InteractionEntity(String interactionId, List<ChatMessage> messages) {
        this.interactionId = interactionId;
        this.messages = messages;
    }

    public String getInteractionId() {
        return interactionId;
    }

    public void setInteractionId(String interactionId) {
        this.interactionId = interactionId;
    }

    public List<ChatMessage> getMessages() {
        return messages;
    }

    public void setMessages(List<ChatMessage> messages) {
        this.messages = messages;
    }

    @Override
    public boolean equals(Object o) {
        if (this == o) return true;
        if (o == null || getClass() != o.getClass()) return false;
        InteractionEntity that = (InteractionEntity) o;
        return Objects.equals(interactionId, that.interactionId);
    }

    @Override
    public int hashCode() {
        return Objects.hash(interactionId, messages);
    }
}

Podemos agora criar o nosso repositório.

package <seupacote>;

import dev.langchain4j.data.message.ChatMessage;
import io.quarkus.mongodb.panache.PanacheMongoRepositoryBase;

import java.util.List;

public class InteractionRepository implements PanacheMongoRepositoryBase<InteractionEntity, String> {

    public InteractionEntity findByInteractionId(String interactionId) {
        return findById(interactionId);
    }

    public void updateMessages(String interactionId, List<ChatMessage> messages) {
        persistOrUpdate(new InteractionEntity(interactionId, messages));
    }

    public void deleteMessages(String interactionId) {
        deleteById(interactionId);
    }

}

Agora iremos implementar alguns componentes do langchain4j, o ChatMemoryStore e o ChatMemoryProvider. O ChatMemoryProvider é a classe que utilizaremos no nosso Agent. Nele iremos adicionar uma ChatMemoryStore que irá utilizar nosso repositório para armazenar as mensagens no nosso MongoDB. Segue o ChatMemoryStore:

package <seupacote>;

import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;

import java.util.List;
import java.util.Objects;

public class MongoDBChatMemoryStore implements ChatMemoryStore {

    private InteractionRepository interactionRepository = new InteractionRepository();

    @Override
    public List<ChatMessage> getMessages(Object memoryId) {
        var interactionEntity = interactionRepository.findByInteractionId(memoryId.toString());
        return Objects.isNull(interactionEntity) ? List.of() : interactionEntity.getMessages();
    }

    @Override
    public void updateMessages(Object memoryId, List<ChatMessage> messages) {
        interactionRepository.updateMessages(memoryId.toString(), messages);
    }

    @Override
    public void deleteMessages(Object memoryId) {
        interactionRepository.deleteMessages(memoryId.toString());
    }
}

O ChatMemoryProvider ficará desse jeito:

package <seupacote>;

import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;

import java.util.function.Supplier;

public class MongoDBChatMemoryProvider implements Supplier<ChatMemoryProvider> {

    private MongoDBChatMemoryStore mongoDBChatMemoryStore = new MongoDBChatMemoryStore();

    @Override
    public ChatMemoryProvider get() {
        return memoryId -> MessageWindowChatMemory.builder()
                .maxMessages(100)
                .id(memoryId)
                .chatMemoryStore(mongoDBChatMemoryStore)
                .build();
    }
}

Repare no MessageWindowChatMemory. É nele que implementamos a janela de mensagens que mencionamos no começo do artigo. No método maxMessages(), você deve alterar pro número que achar melhor pro seu cenário. O que recomendo é utilizar o maior número de mensagens que já existiu no seu cenário, ou utilizar a média. Aqui definimos o número arbitrário 100.

Vamos agora alterar o nosso agente para utilizar o nosso ChatMemoryProvider novo e adicionar MemoryId. Ele deve ficar assim:

package <seupacote>;

import dev.langchain4j.service.SystemMessage;
import dev.langchain4j.service.UserMessage;
import io.quarkiverse.langchain4j.RegisterAiService;
import io.quarkiverse.langchain4j.ToolBox;
import jakarta.enterprise.context.ApplicationScoped;

@ApplicationScoped
@RegisterAiService(
        chatMemoryProviderSupplier = MongoDBChatMemoryProvider.class
)
public interface Agent {

    @ToolBox(AgentTools.class)
    @SystemMessage("""
            Você é um agente especializado em futebol brasileiro, seu nome é FutAgentBR
            Você sabe responder sobre os principais títulos dos principais times brasileiros e da seleção brasileira
            Sua resposta precisa ser educada, você pode deve responder em Português brasileiro e de forma relevante a pergunta feita

            Quando você não souber a resposta, responda que você não sabe responder nesse momento mas saberá em futuras versões.
            """)
    String chat(@MemoryId String interactionId, @UserMessage String message);
}

Isso deve quebrar o nosso AgentWSEndpoint. Vamos alterá-lo para que ele receba o identificador da Interaction e possamos utilizar como nosso MemoryId:

package <seupacote>;

import io.quarkus.websockets.next.OnTextMessage;
import io.quarkus.websockets.next.WebSocket;
import io.quarkus.websockets.next.WebSocketConnection;
import jakarta.inject.Inject;

import java.util.Objects;
import java.util.UUID;

@WebSocket(path = "/ws/{interactionId}")
public class AgentWSEndpoint {

    private final Agent agent;

    private final WebSocketConnection connection;

    @Inject
    AgentWSEndpoint(Agent agent, WebSocketConnection connection) {
        this.agent = agent;
        this.connection = connection;
    }

    @OnTextMessage
    String reply(String message) {
        var interactionId = connection.pathParam("interactionId");
        return agent.chat(
                Objects.isNull(interactionId) || interactionId.isBlank()
                        ? UUID.randomUUID().toString()
                        : interactionId,
                message
        );
    }

}

Já podemos testar o nosso agente novamente. Para isso basta conectar-mos no websocket passando um UUID sempre que quisermos. Você pode gerar um novo UUID aqui, ou utilizar o comando uuidgen no linux.

Ao realizarmos o teste você não receberá resposta alguma do agente. Isso acontece por quê o agente está tendo problemas ao gravar nossas mensagens no MongoDB e ele te mostrará isso através de uma exceção. Para que possamos verificar essa exceção acontecendo, devemos incluir uma nova propriedade no nosso src/main/resources/application.properties, que é o nível do log que queremos ver no Quarkus. Então, adicione a seguinte linha nele:

quarkus.log.level=DEBUG

Agora teste o agente. A exceção deve ser essa:

DEBUG [io.qua.web.nex.run.Endpoints] (vert.x-eventloop-thread-1) Connection closed due to unhandled failure org.bson.codecs.configuration.CodecConfigurationException: An exception occurred when encoding using the AutomaticPojoCodec.

Essa exceção ocorre porque o MongoDB não consegue lidar com a interface ChatMessage do Langchain4j, então devemos implementar um codec pra que isso seja possível. O próprio Quarkus já nos oferece um codec, mas precisamos deixar explicito que queremos utilizar ele. Criaremos então as classes ChatMessageCodec e ChatMessageCodecProvider como segue:

package <seupacote>;

import com.mongodb.MongoClientSettings;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ChatMessageJsonCodec;
import io.quarkiverse.langchain4j.QuarkusChatMessageJsonCodecFactory;
import org.bson.BsonReader;
import org.bson.BsonValue;
import org.bson.BsonWriter;
import org.bson.Document;
import org.bson.codecs.Codec;
import org.bson.codecs.CollectibleCodec;
import org.bson.codecs.DecoderContext;
import org.bson.codecs.EncoderContext;

public class ChatMessageCodec implements CollectibleCodec<ChatMessage> {

    private final Codec<Document> documentCodec;
    private final ChatMessageJsonCodec chatMessageJsonCodec;

    public ChatMessageCodec() {
        this.documentCodec = MongoClientSettings.getDefaultCodecRegistry().get(Document.class);
        this.chatMessageJsonCodec = new QuarkusChatMessageJsonCodecFactory().create();
    }

    @Override
    public ChatMessage generateIdIfAbsentFromDocument(ChatMessage document) {
        return document;
    }

    @Override
    public boolean documentHasId(ChatMessage document) {
        return false;
    }

    @Override
    public BsonValue getDocumentId(ChatMessage document) {
        return null;
    }

    @Override
    public ChatMessage decode(BsonReader reader, DecoderContext decoderContext) {
        var document = documentCodec.decode(reader, decoderContext);
        return this.chatMessageJsonCodec.messageFromJson(document.toJson());
    }

    @Override
    public void encode(BsonWriter writer, ChatMessage value, EncoderContext encoderContext) {
        var json = this.chatMessageJsonCodec.messageToJson(value);
        var doc = Document.parse(json);
        documentCodec.encode(writer, doc, encoderContext);
    }

    @Override
    public Class<ChatMessage> getEncoderClass() {
        return ChatMessage.class;
    }
}

package <seupacote>;

import dev.langchain4j.data.message.ChatMessage;
import org.bson.codecs.Codec;
import org.bson.codecs.configuration.CodecProvider;
import org.bson.codecs.configuration.CodecRegistry;

public class ChatMessageCodecProvider implements CodecProvider {

    @Override
    public <T> Codec<T> get(Class<T> clazz, CodecRegistry registry) {
        if (clazz == ChatMessage.class) {
            return (Codec<T>) new ChatMessageCodec();
        }
        return null;
    }

}

Pronto! Agora podemos testar e verificar as mensagens no nosso MongoDB. Ao consultarmos, podemos verificar os 3 tipos de mensagens no array messages do documento.

Isso encerra a segunda parte da nossa série. Esperamos que tenham gostado e até a parte 3.