[Fix #836] Implement beforeAgentInvocation and afterAgentInvocation f…

experimental/fluent/agentic-langchain4j/src/main/java/io/serverlessworkflow/fluent/agentic/langchain4j/ConditionalAgentServiceImpl.java

@@ -20,15 +20,23 @@
 import dev.langchain4j.agentic.internal.AgentExecutor;
 import dev.langchain4j.agentic.scope.AgenticScope;
 import dev.langchain4j.agentic.workflow.ConditionalAgentService;
+import io.serverlessworkflow.fluent.agentic.AgenticScopedRequest;
+import io.serverlessworkflow.fluent.agentic.AgenticScopedResponse;
 import java.util.Arrays;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 import java.util.function.Predicate;
 
 public class ConditionalAgentServiceImpl<T>
     extends AbstractAgentService<T, ConditionalAgentService<T>>
     implements ConditionalAgentService<T> {
 
+  private final AtomicReference<Consumer<AgenticScopedRequest>> beforeAgentInvocation =
+      new AtomicReference<>();
+  private final AtomicReference<Consumer<AgenticScopedResponse>> afterAgentInvocation =
+      new AtomicReference<>();
+
   private ConditionalAgentServiceImpl(Class<T> agentServiceClass) {
     super(agentServiceClass);
   }
@@ -39,7 +47,21 @@ public static <T> ConditionalAgentService<T> builder(Class<T> agentServiceClass)
 
   @Override
   public ConditionalAgentService<T> subAgents(Object... agents) {
-    this.workflowBuilder.tasks(t -> t.sequence(agents));
+    this.workflowBuilder.tasks(
+        t ->
+            t.sequence(agents)
+                .inputFrom(
+                    scope -> {
+                      if (beforeAgentInvocation.get() != null) {
+                        beforeAgentInvocation.get().accept((AgenticScopedRequest) scope);
+                      }
+                    })
+                .outputAs(
+                    scope -> {
+                      if (afterAgentInvocation.get() != null) {
+                        afterAgentInvocation.get().accept((AgenticScopedResponse) scope);
+                      }
+                    }));
     return this;
   }
 
@@ -50,20 +72,43 @@ public ConditionalAgentService<T> subAgents(List<AgentExecutor> agentExecutors)
 
   @Override
   public ConditionalAgentService<T> beforeAgentInvocation(Consumer<AgentRequest> consumer) {
-    throw new UnsupportedOperationException(
-        "Feature not implemented yet. See: https://github.com/serverlessworkflow/sdk-java/issues/836");
+    beforeAgentInvocation.set(
+        agenticScopedRequest -> consumer.accept(agenticScopedRequest.asAgentRequest()));
+    return this;
   }
 
   @Override
   public ConditionalAgentService<T> afterAgentInvocation(Consumer<AgentResponse> consumer) {
-    throw new UnsupportedOperationException(
-        "Feature not implemented yet. See: https://github.com/serverlessworkflow/sdk-java/issues/836");
+    afterAgentInvocation.set(
+        agenticScopedResponse -> consumer.accept(agenticScopedResponse.asAgentResponse()));
+    return this;
   }
 
   @Override
   public ConditionalAgentService<T> subAgents(Predicate<AgenticScope> condition, Object... agents) {
     this.workflowBuilder.tasks(
-        t -> Arrays.stream(agents).forEach(agent -> t.when(condition).agent(agent)));
+        t ->
+            Arrays.stream(agents)
+                .forEach(
+                    agent ->
+                        t.when(condition)
+                            .agent(agent)
+                            .inputFrom(
+                                scope -> {
+                                  if (beforeAgentInvocation.get() != null) {
+                                    beforeAgentInvocation
+                                        .get()
+                                        .accept((AgenticScopedRequest) scope);
+                                  }
+                                })
+                            .outputAs(
+                                scope -> {
+                                  if (afterAgentInvocation.get() != null) {
+                                    afterAgentInvocation
+                                        .get()
+                                        .accept((AgenticScopedResponse) scope);
+                                  }
+                                })));
     return this;
   }
 
@@ -76,7 +121,22 @@ public ConditionalAgentService<T> subAgents(
   @Override
   public ConditionalAgentService<T> subAgent(
       Predicate<AgenticScope> condition, AgentExecutor agentExecutor) {
-    this.workflowBuilder.tasks(t -> t.when(condition).agent(agentExecutor));
+    this.workflowBuilder.tasks(
+        t ->
+            t.when(condition)
+                .agent(agentExecutor)
+                .inputFrom(
+                    scope -> {
+                      if (beforeAgentInvocation.get() != null) {
+                        beforeAgentInvocation.get().accept((AgenticScopedRequest) scope);
+                      }
+                    })
+                .outputAs(
+                    scope -> {
+                      if (afterAgentInvocation.get() != null) {
+                        afterAgentInvocation.get().accept((AgenticScopedResponse) scope);
+                      }
+                    }));
     return this;
   }
 }

experimental/fluent/agentic/src/main/java/io/serverlessworkflow/fluent/agentic/AgentAdapters.java

@@ -22,7 +22,9 @@
 import dev.langchain4j.agentic.scope.DefaultAgenticScope;
 import io.serverlessworkflow.api.types.func.LoopPredicateIndex;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.BiPredicate;
+import java.util.function.Consumer;
 import java.util.function.Function;
 import java.util.function.Predicate;
 
@@ -34,8 +36,26 @@ public static List<AgentExecutor> toExecutors(Object... agents) {
     return agentsToExecutors(agents);
   }
 
-  public static Function<DefaultAgenticScope, Object> toFunction(AgentExecutor exec) {
-    return exec::execute;
+  public static Function<DefaultAgenticScope, Object> toFunction(
+      AgentExecutor exec,
+      AtomicReference<Consumer<AgenticScopedRequest>> beforeAgentInvocation,
+      AtomicReference<Consumer<AgenticScopedResponse>> afterAgentInvocation) {
+    String agentName = exec.agentInvoker().name();
+    return defaultAgenticScope -> {
+      if (beforeAgentInvocation.get() != null) {
+        beforeAgentInvocation
+            .get()
+            .accept(new AgenticScopedRequest(defaultAgenticScope, agentName));
+      }
+      Object result = exec.execute(defaultAgenticScope);
+      if (afterAgentInvocation.get() != null) {
+        defaultAgenticScope.writeState("input", result);
+        afterAgentInvocation
+            .get()
+            .accept(new AgenticScopedResponse(defaultAgenticScope, agentName, result));
+      }
+      return result;
+    };
   }
 
   public static LoopPredicateIndex<AgenticScope, Object> toWhile(Predicate<AgenticScope> exit) {

experimental/fluent/agentic/src/main/java/io/serverlessworkflow/fluent/agentic/AgentDoTaskBuilder.java

@@ -15,6 +15,7 @@
  */
 package io.serverlessworkflow.fluent.agentic;
 
+import dev.langchain4j.agentic.scope.AgenticScope;
 import io.serverlessworkflow.fluent.agentic.spi.AgentDoFluent;
 import io.serverlessworkflow.fluent.func.FuncCallTaskBuilder;
 import io.serverlessworkflow.fluent.func.FuncEmitTaskBuilder;
@@ -117,4 +118,14 @@ public AgentDoTaskBuilder switchCase(
     this.listBuilder().switchCase(name, itemsConfigurer);
     return self();
   }
+
+  public AgentDoTaskBuilder inputFrom(Consumer<AgenticScope> inputFrom) {
+    this.listBuilder().inputFrom(inputFrom);
+    return self();
+  }
+
+  public AgentDoTaskBuilder outputAs(Consumer<AgenticScope> outputAs) {
+    this.listBuilder().outputAs(outputAs);
+    return self();
+  }
 }

experimental/fluent/agentic/src/main/java/io/serverlessworkflow/fluent/agentic/AgentTaskItemListBuilder.java

@@ -16,6 +16,7 @@
 package io.serverlessworkflow.fluent.agentic;
 
 import dev.langchain4j.agentic.internal.AgentExecutor;
+import dev.langchain4j.agentic.scope.AgenticScope;
 import dev.langchain4j.agentic.scope.DefaultAgenticScope;
 import io.serverlessworkflow.api.types.Task;
 import io.serverlessworkflow.api.types.TaskItem;
@@ -29,12 +30,17 @@
 import io.serverlessworkflow.fluent.func.FuncTaskItemListBuilder;
 import io.serverlessworkflow.fluent.spec.BaseTaskItemListBuilder;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Consumer;
 
 public class AgentTaskItemListBuilder extends BaseTaskItemListBuilder<AgentTaskItemListBuilder>
     implements AgentDoFluent<AgentTaskItemListBuilder> {
 
   private final FuncTaskItemListBuilder delegate;
+  private final AtomicReference<Consumer<AgenticScopedRequest>> beforeAgentInvocation =
+      new AtomicReference<>();
+  private final AtomicReference<Consumer<AgenticScopedResponse>> afterAgentInvocation =
+      new AtomicReference<>();
 
   public AgentTaskItemListBuilder() {
     super();
@@ -58,7 +64,11 @@ public AgentTaskItemListBuilder agent(String name, Object agent) {
             exec ->
                 this.delegate.function(
                     name,
-                    fn -> fn.function(AgentAdapters.toFunction(exec), DefaultAgenticScope.class)));
+                    fn ->
+                        fn.function(
+                            AgentAdapters.toFunction(
+                                exec, beforeAgentInvocation, afterAgentInvocation),
+                            DefaultAgenticScope.class)));
     return self();
   }
 
@@ -96,7 +106,10 @@ public AgentTaskItemListBuilder parallel(String name, Object... agents) {
                 ex.agentInvoker().name() != null
                     ? ex.agentInvoker().name()
                     : "branch-" + i + "-" + name;
-            fork.branch(agentName, AgentAdapters.toFunction(ex), DefaultAgenticScope.class);
+            fork.branch(
+                agentName,
+                AgentAdapters.toFunction(ex, beforeAgentInvocation, afterAgentInvocation),
+                DefaultAgenticScope.class);
           }
         });
     return self();
@@ -154,4 +167,14 @@ public AgentTaskItemListBuilder switchCase(
     this.delegate.switchCase(name, itemsConfigurer);
     return self();
   }
+
+  public AgentTaskItemListBuilder inputFrom(Consumer<AgenticScope> beforeAgentInvocation) {
+    this.beforeAgentInvocation.set(beforeAgentInvocation::accept);
+    return self();
+  }
+
+  public AgentTaskItemListBuilder outputAs(Consumer<AgenticScope> afterAgentInvocationConsumer) {
+    this.afterAgentInvocation.set(afterAgentInvocationConsumer::accept);
+    return self();
+  }
 }

experimental/fluent/agentic/src/main/java/io/serverlessworkflow/fluent/agentic/AgenticScopedRequest.java

@@ -0,0 +1,87 @@
+/*
+ * Copyright 2020-Present The Serverless Workflow Specification Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.serverlessworkflow.fluent.agentic;
+
+import dev.langchain4j.agentic.agent.AgentRequest;
+import dev.langchain4j.agentic.internal.AgentInvocation;
+import dev.langchain4j.agentic.scope.AgenticScope;
+import java.util.List;
+import java.util.Map;
+
+public class AgenticScopedRequest implements AgenticScope {
+
+  protected final String agentName;
+  private final AgenticScope wrapped;
+
+  AgenticScopedRequest(AgenticScope wrapped, String agentName) {
+    this.wrapped = wrapped;
+    this.agentName = agentName;
+  }
+
+  @Override
+  public Object memoryId() {
+    return wrapped.memoryId();
+  }
+
+  @Override
+  public void writeState(String s, Object o) {
+    wrapped.writeState(s, o);
+  }
+
+  @Override
+  public void writeStates(Map<String, Object> map) {
+    wrapped.writeStates(map);
+  }
+
+  @Override
+  public boolean hasState(String s) {
+    return wrapped.hasState(s);
+  }
+
+  @Override
+  public Object readState(String s) {
+    return wrapped.readState(s);
+  }
+
+  @Override
+  public <T> T readState(String s, T t) {
+    return wrapped.readState(s, t);
+  }
+
+  @Override
+  public Map<String, Object> state() {
+    return wrapped.state();
+  }
+
+  @Override
+  public String contextAsConversation(String... strings) {
+    return wrapped.contextAsConversation(strings);
+  }
+
+  @Override
+  public String contextAsConversation(Object... objects) {
+    return wrapped.contextAsConversation(objects);
+  }
+
+  @Override
+  public List<AgentInvocation> agentInvocations(String s) {
+    return wrapped.agentInvocations(s);
+  }
+
+  public AgentRequest asAgentRequest() {
+    return new AgentRequest(this, agentName, state());
+  }
+}

experimental/fluent/agentic/src/main/java/io/serverlessworkflow/fluent/agentic/AgenticScopedResponse.java

@@ -0,0 +1,33 @@
+/*
+ * Copyright 2020-Present The Serverless Workflow Specification Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package io.serverlessworkflow.fluent.agentic;
+
+import dev.langchain4j.agentic.agent.AgentResponse;
+import dev.langchain4j.agentic.scope.AgenticScope;
+
+public class AgenticScopedResponse extends AgenticScopedRequest {
+
+  private final Object response;
+
+  AgenticScopedResponse(AgenticScope wrapped, String agentName, Object response) {
+    super(wrapped, agentName);
+    this.response = response;
+  }
+
+  public AgentResponse asAgentResponse() {
+    return new AgentResponse(this, agentName, state(), response);
+  }
+}