fix: Agent extend conversation with assistant's reply

Author: bsorrentino

agents/src/main/java/dev/langchain4j/Agent.java

@@ -6,6 +6,7 @@
 import dev.langchain4j.data.message.SystemMessage;
 import dev.langchain4j.data.message.UserMessage;
 import dev.langchain4j.data.message.ToolExecutionResultMessage;
+
 import dev.langchain4j.model.chat.ChatLanguageModel;
 import dev.langchain4j.model.input.PromptTemplate;
 import dev.langchain4j.model.output.Response;
@@ -16,53 +17,39 @@
 import java.util.List;
 import java.util.Map;
 
+import static java.lang.String.format;
+
 @Builder
 public class Agent {
 
     private final  ChatLanguageModel chatLanguageModel;
     @Singular  private final List<ToolSpecification> tools;
 
 
-    public Response<AiMessage> execute( Map<String,Object> inputs ) {
-        var messages = new ArrayList<ChatMessage>();
-        var promptTemplate = PromptTemplate.from( "USER: {{input}}" ).apply(inputs);
-
-        messages.add(new SystemMessage("You are a helpful assistant"));
-
-        messages.add( new UserMessage(promptTemplate.text())  );
-
-        return chatLanguageModel.generate( messages, tools );
-    }
-
-    private PromptTemplate getToolResponseTemplate( ) {
-        var TEMPLATE_TOOL_RESPONSE = """
-                TOOL RESPONSE:
-                ---------------------
-                {{observation}}
-                --------------------
-                """;
-        return PromptTemplate.from(TEMPLATE_TOOL_RESPONSE);
-    }
-
     public Response<AiMessage> execute( String input, List<AgentExecutor.IntermediateStep> intermediateSteps ) {
-        var agentScratchpadTemplate = getToolResponseTemplate();
-        var userMessageTemplate = PromptTemplate.from( "USER'S INPUT: {{input}}" )
+        var userMessageTemplate = PromptTemplate.from( "{{input}}" )
                                                     .apply( Map.of( "input", input));
 
         var messages = new ArrayList<ChatMessage>();
 
         messages.add(new SystemMessage("You are a helpful assistant"));
+        messages.add(new UserMessage(userMessageTemplate.text()));
 
-        if( intermediateSteps.isEmpty()) {
-            messages.add(new UserMessage(userMessageTemplate.text()));
-        }
+        if (!intermediateSteps.isEmpty()) {
 
-        for( AgentExecutor.IntermediateStep step: intermediateSteps ) {
-            var agentScratchpad = agentScratchpadTemplate
-                                        .apply( Map.of("observation", step.observation()) );
-            messages.add(new UserMessage(agentScratchpad.text()));
-        }
+            var toolRequests = intermediateSteps.stream()
+                    .map(AgentExecutor.IntermediateStep::action)
+                    .map(AgentExecutor.AgentAction::toolExecutionRequest)
+                    .toList();
+
+            messages.add(new AiMessage(toolRequests)); // reply with tool requests
 
+            for (AgentExecutor.IntermediateStep step : intermediateSteps) {
+                var toolRequest = step.action().toolExecutionRequest();
+
+                messages.add(new ToolExecutionResultMessage(toolRequest.id(), toolRequest.name(), step.observation()));
+            }
+        }
         return chatLanguageModel.generate( messages, tools );
     }
 }

agents/src/main/java/dev/langchain4j/AgentExecutor.java

@@ -2,6 +2,7 @@
 
 import dev.langchain4j.agent.tool.ToolExecutionRequest;
 import dev.langchain4j.agent.tool.ToolSpecification;
+import dev.langchain4j.model.chat.ChatLanguageModel;
 import dev.langchain4j.model.openai.OpenAiChatModel;
 import dev.langchain4j.model.output.FinishReason;
 import org.bsc.langgraph4j.GraphState;
@@ -117,21 +118,9 @@ String shouldContinue(State state) {
         return "continue";
     }
 
-    public AsyncIterator<GraphState.Runnable.NodeOutput<State>> execute(Map<String, Object> inputs, List<Object> objectsWithTools) throws Exception {
+    public AsyncIterator<GraphState.Runnable.NodeOutput<State>> execute(ChatLanguageModel chatLanguageModel, Map<String, Object> inputs, List<Object> objectsWithTools) throws Exception {
 
 
-        var openApiKey = DotEnvConfig.valueOf("OPENAI_API_KEY")
-                            .orElseThrow( () -> new IllegalArgumentException("no APIKEY provided!"));
-
-        var chatLanguageModel = OpenAiChatModel.builder()
-                .apiKey( openApiKey )
-                .modelName( "gpt-3.5-turbo-0613" )
-                .logResponses(true)
-                .maxRetries(2)
-                .temperature(0.0)
-                .maxTokens(2000)
-                .build();
-
         var toolInfoList = ToolInfo.fromList( objectsWithTools );
 
         final List<ToolSpecification> toolSpecifications = toolInfoList.stream()

agents/src/test/java/dev/langchain4j/AgentExecutorTest.java

@@ -2,6 +2,7 @@
 
 import dev.langchain4j.agent.tool.P;
 import dev.langchain4j.agent.tool.Tool;
+import dev.langchain4j.model.openai.OpenAiChatModel;
 
 import java.util.List;
 import java.util.Map;
@@ -15,11 +16,25 @@ public static void main( String[] args)  {
 
         DotEnvConfig.load();
 
-        var agentExecutor = new AgentExecutor();
+        var openApiKey = DotEnvConfig.valueOf("OPENAI_API_KEY")
+                .orElseThrow( () -> new IllegalArgumentException("no APIKEY provided!"));
+
+        var chatLanguageModel = OpenAiChatModel.builder()
+                .apiKey( openApiKey )
+                .modelName( "gpt-3.5-turbo-0125" )
+                .logResponses(true)
+                .maxRetries(2)
+                .temperature(0.0)
+                .maxTokens(2000)
+                .build();
 
         try {
+            var agentExecutor = new AgentExecutor();
+
             var iterator = agentExecutor.execute(
-                    Map.of( "input", "what is the result of test with message: 'MY FIRST TEST'?"),
+                    chatLanguageModel,
+                    //Map.of( "input", "what is the result of test with messages: 'MY FIRST TEST' and the result of test with message: 'MY SECOND TEST'"),
+                    Map.of( "input", "what is the result of test with messages: 'MY FIRST TEST'"),
                     List.of(new TestTool()) );
 
            AgentExecutor.State output = null;

agents/src/test/java/dev/langchain4j/AgentTest.java

@@ -6,6 +6,7 @@
 import org.junit.jupiter.api.BeforeAll;
 import dev.langchain4j.model.openai.OpenAiChatModel;
 
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 
@@ -40,7 +41,7 @@ public static void main( String[] args) throws Exception  {
                 .build();
 
         var msg = "hello world";
-        var response = agent.execute( Map.of("input", format("this is an AI test with message: '%s'", msg) ));
+        var response = agent.execute( format("this is an AI test with message: '%s'", msg), List.of() );
 
         assertNotNull(response);
         assertEquals(response.finishReason(), FinishReason.TOOL_EXECUTION );