Fix: Persist ChatMemory when Tool ReturnBehavior is IMMEDIATE (#2017)

[codesmith25103]

• Fixes: Bug: ChatMemoryStore.updateMessages() is never called when a tool with ReturnBehavior.IMMEDIATE is executed #2017

[geoand]

It would be awesome if we could have a test for this

[codesmith25103]

Hey @geoand I was trying to write test case but facing some difficulty in doing so, can you guide me little so I can successfully implement it.

[codesmith25103]

package org.acme.example.openai;

import static org.junit.jupiter.api.Assertions.assertTrue;

import java.util.ArrayList;
import java.util.List;

import jakarta.annotation.Priority;

import jakarta.enterprise.context.ApplicationScoped;
import jakarta.enterprise.inject.Alternative;
import jakarta.enterprise.inject.Produces;
import jakarta.inject.Inject;

import org.junit.jupiter.api.Test;

import dev.langchain4j.agent.tool.ReturnBehavior;
import dev.langchain4j.agent.tool.Tool;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.service.Result;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import io.quarkiverse.langchain4j.RegisterAiService;
import io.quarkus.test.junit.QuarkusTest;

@QuarkusTest
public class ImmediateReturnMemoryTest {

    @Inject
    MyAiService aiService;

    @Inject
    SpyMemoryStore spyStore;

    @Test
    public void testImmediateReturnPersistsToMemory() {

        try {
            aiService.chat("Please run the tool my_tool");
        } catch (Exception e) {

        }

        assertTrue(spyStore.updateCalled,
                "FAILURE: ChatMemoryStore.updateMessages() was never called! The tool result was lost.");
    }

    @RegisterAiService
    public interface MyAiService {
        Result<String> chat(String userMessage);
    }

    @ApplicationScoped
    public static class MyTool {
        @Tool(name = "my_tool", returnBehavior = ReturnBehavior.IMMEDIATE)
        public String doWork(String input) {
            return "TOOL_EXECUTED_IMMEDIATELY";
        }
    }

    @ApplicationScoped
    public static class SpyMemoryStore implements ChatMemoryStore {
        public boolean updateCalled = false;

        @Override
        public List<ChatMessage> getMessages(Object memoryId) {
            return new ArrayList<>();
        }

        @Override
        public void updateMessages(Object memoryId, List<ChatMessage> messages) {
            this.updateCalled = true;
        }

        @Override
        public void deleteMessages(Object memoryId) {
        }
    }

    public static class MemoryConfig {

        @Produces
        @ApplicationScoped
        @Alternative
        @Priority(1)
        public ChatMemoryProvider memoryProvider(SpyMemoryStore store) {
            return memoryId -> MessageWindowChatMemory.builder()
                    .maxMessages(10)
                    .chatMemoryStore(store)
                    .id(memoryId)
                    .build();
        }
    }
}

[codesmith25103]

However, when I ran the test case after building the project, it still failed.

[geoand]

Did you try debugging it?

[stanlayze]

Should we commit the messages after the type-check in the next 3 lines?

[geoand]

Yeah, that probably makes sense

[stanlayze]

Hi @codesmith25103,
I was playing around with the test. Since I am not so experienced in this project I had a look around existing tests. I found some useful test setups in the io.quarkiverse.langchain4j.test.tools package of the core-module. Being inspired by them I came up with this test:

package io.quarkiverse.langchain4j.test.tools;

import dev.langchain4j.agent.tool.ReturnBehavior;
import dev.langchain4j.agent.tool.Tool;
import dev.langchain4j.agent.tool.ToolExecutionRequest;
import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.ToolExecutionResultMessage;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.ChatMemoryProvider;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.chat.ChatModel;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.output.FinishReason;
import dev.langchain4j.model.output.TokenUsage;
import dev.langchain4j.service.MemoryId;
import dev.langchain4j.service.Result;
import dev.langchain4j.service.UserMessage;
import dev.langchain4j.store.memory.chat.ChatMemoryStore;
import io.quarkiverse.langchain4j.RegisterAiService;
import io.quarkiverse.langchain4j.ToolBox;
import io.quarkiverse.langchain4j.test.Lists;
import io.quarkus.test.QuarkusUnitTest;
import jakarta.enterprise.context.control.ActivateRequestContext;
import jakarta.inject.Inject;
import jakarta.inject.Singleton;
import org.assertj.core.api.Assertions;
import org.jboss.shrinkwrap.api.ShrinkWrap;
import org.jboss.shrinkwrap.api.spec.JavaArchive;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.api.extension.RegisterExtension;

import java.util.ArrayList;
import java.util.List;
import java.util.UUID;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Supplier;

class ToolExecutionImmediateReturnTest {

    @RegisterExtension
    static final QuarkusUnitTest unitTest = new QuarkusUnitTest()
            .setArchiveProducer(() -> ShrinkWrap.create(JavaArchive.class).addClasses(MyAiService.class, Lists.class));

    @Inject
    MyAiService aiService;

    @Test
    @ActivateRequestContext
    void testImmediateReturn() {
        String uuid = UUID.randomUUID().toString();
        aiService.hello("abc", "hiImmediate - " + uuid);

        Assertions.assertThat(MyMemoryProviderSupplier.spystore.updateInvocationCounter).hasValueGreaterThan(0);
        Assertions.assertThat(MyMemoryProviderSupplier.spystore.messages).anyMatch(ToolExecutionResultMessage.class::isInstance, "Tool execution result message should be present");
    }

    @RegisterAiService(chatLanguageModelSupplier = MyChatModelSupplier.class, chatMemoryProviderSupplier = MyMemoryProviderSupplier.class)
    public interface MyAiService {

        @ToolBox(MyTool.class)
        Result<String> hello(@MemoryId String memoryId, @UserMessage String userMessageContainingTheToolId);

    }

    @Singleton
    public static class MyTool {
        @Tool(returnBehavior = ReturnBehavior.IMMEDIATE)
        public String hiImmediate(String m) {
            System.out.println("Executing tool with: " + m);
            return "hiImmediate";
        }

    }

    public static class MyChatModelSupplier implements Supplier<ChatModel> {

        @Override
        public ChatModel get() {
            return new MyChatModel();
        }
    }

    public static class MyChatModel implements ChatModel {

        @Override
        public ChatResponse chat(List<ChatMessage> messages) {
            throw new UnsupportedOperationException("Should not be called");
        }

        @Override
        public ChatResponse doChat(ChatRequest chatRequest) {
            List<ChatMessage> messages = chatRequest.messages();
            if (messages.size() == 1) {
                // Only the user message, extract the tool id from it
                String text = ((dev.langchain4j.data.message.UserMessage) messages.get(0)).singleText();
                var segments = text.split(" - ");
                var toolId = segments[0];
                var content = segments[1];
                // Only the user message
                return ChatResponse.builder().aiMessage(new AiMessage("cannot be blank", List.of(ToolExecutionRequest.builder()
                        .id("my-tool-" + toolId)
                        .name(toolId)
                        .arguments("{\"m\":\"" + content + "\"}")
                        .build()))).tokenUsage(new TokenUsage(0, 0)).finishReason(FinishReason.TOOL_EXECUTION).build();
            } else if (messages.size() == 3) {
                // user -> tool request -> tool response
                ToolExecutionResultMessage last = (ToolExecutionResultMessage) Lists.last(messages);
                return ChatResponse.builder().aiMessage(AiMessage.from("response: " + last.text())).build();

            }
            return ChatResponse.builder().aiMessage(new AiMessage("Unexpected")).build();
        }
    }

    public static class SpyMemoryStore implements ChatMemoryStore {
        public AtomicInteger updateInvocationCounter = new AtomicInteger(0);
        public List<ChatMessage> messages = new ArrayList<>();

        @Override
        public List<ChatMessage> getMessages(Object memoryId) {
            return messages;
        }

        @Override
        public void updateMessages(Object memoryId, List<ChatMessage> messages) {
            this.messages.addAll(messages);
            this.updateInvocationCounter.incrementAndGet();
        }

        @Override
        public void deleteMessages(Object memoryId) {
        }
    }

    public static class MyMemoryProviderSupplier implements Supplier<ChatMemoryProvider> {

        static SpyMemoryStore spystore = new SpyMemoryStore();

        @Override
        public ChatMemoryProvider get() {
            return new ChatMemoryProvider() {
                @Override
                public ChatMemory get(Object memoryId) {
                    return MessageWindowChatMemory.builder()
                            .maxMessages(10)
                            .chatMemoryStore(spystore)
                            .id(memoryId)
                            .build();
                }
            };
        }
    }
}

Do you want to try it and update your PR? Or should I create a new fork and PR?
fyi @geoand

[codesmith25103]

@stanlayze thanks for your help. Yes , I will try and update my PR ASAP.
Thank you