Skip to content

Commit 7b5ebfa

Browse files
treblereeltreblereel
committed
[Fix #836] Implement beforeAgentInvocation and afterAgentInvocation for LC4J
Signed-off-by: Dmitrii Tikhomirov <[email protected]>
1 parent 9a16384 commit 7b5ebfa

File tree

5 files changed

+276
-5
lines changed

5 files changed

+276
-5
lines changed

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

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,9 @@
2222
import dev.langchain4j.agentic.scope.DefaultAgenticScope;
2323
import io.serverlessworkflow.api.types.func.LoopPredicateIndex;
2424
import java.util.List;
25+
import java.util.concurrent.atomic.AtomicReference;
2526
import java.util.function.BiPredicate;
27+
import java.util.function.Consumer;
2628
import java.util.function.Function;
2729
import java.util.function.Predicate;
2830

@@ -34,8 +36,21 @@ public static List<AgentExecutor> toExecutors(Object... agents) {
3436
return agentsToExecutors(agents);
3537
}
3638

37-
public static Function<DefaultAgenticScope, Object> toFunction(AgentExecutor exec) {
38-
return exec::execute;
39+
public static Function<DefaultAgenticScope, Object> toFunction(
40+
AgentExecutor exec,
41+
AtomicReference<Consumer<AgenticScope>> beforeAgentInvocation,
42+
AtomicReference<Consumer<AgenticScope>> afterAgentInvocation) {
43+
return defaultAgenticScope -> {
44+
if (beforeAgentInvocation.get() != null) {
45+
beforeAgentInvocation.get().accept(defaultAgenticScope);
46+
}
47+
Object result = exec.execute(defaultAgenticScope);
48+
if (afterAgentInvocation.get() != null) {
49+
defaultAgenticScope.writeState("input", result);
50+
afterAgentInvocation.get().accept(defaultAgenticScope);
51+
}
52+
return result;
53+
};
3954
}
4055

4156
public static LoopPredicateIndex<AgenticScope, Object> toWhile(Predicate<AgenticScope> exit) {

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

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,7 @@
1515
*/
1616
package io.serverlessworkflow.fluent.agentic;
1717

18+
import dev.langchain4j.agentic.scope.AgenticScope;
1819
import io.serverlessworkflow.fluent.agentic.spi.AgentDoFluent;
1920
import io.serverlessworkflow.fluent.func.FuncCallTaskBuilder;
2021
import io.serverlessworkflow.fluent.func.FuncEmitTaskBuilder;
@@ -117,4 +118,14 @@ public AgentDoTaskBuilder switchCase(
117118
this.listBuilder().switchCase(name, itemsConfigurer);
118119
return self();
119120
}
121+
122+
public AgentDoTaskBuilder inputFrom(Consumer<AgenticScope> inputFrom) {
123+
this.listBuilder().inputFrom(inputFrom);
124+
return self();
125+
}
126+
127+
public AgentDoTaskBuilder outputAs(Consumer<AgenticScope> outputAs) {
128+
this.listBuilder().outputAs(outputAs);
129+
return self();
130+
}
120131
}

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

Lines changed: 25 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
package io.serverlessworkflow.fluent.agentic;
1717

1818
import dev.langchain4j.agentic.internal.AgentExecutor;
19+
import dev.langchain4j.agentic.scope.AgenticScope;
1920
import dev.langchain4j.agentic.scope.DefaultAgenticScope;
2021
import io.serverlessworkflow.api.types.Task;
2122
import io.serverlessworkflow.api.types.TaskItem;
@@ -29,12 +30,17 @@
2930
import io.serverlessworkflow.fluent.func.FuncTaskItemListBuilder;
3031
import io.serverlessworkflow.fluent.spec.BaseTaskItemListBuilder;
3132
import java.util.List;
33+
import java.util.concurrent.atomic.AtomicReference;
3234
import java.util.function.Consumer;
3335

3436
public class AgentTaskItemListBuilder extends BaseTaskItemListBuilder<AgentTaskItemListBuilder>
3537
implements AgentDoFluent<AgentTaskItemListBuilder> {
3638

3739
private final FuncTaskItemListBuilder delegate;
40+
private final AtomicReference<Consumer<AgenticScope>> beforeAgentInvocation =
41+
new AtomicReference<>();
42+
private final AtomicReference<Consumer<AgenticScope>> afterAgentInvocation =
43+
new AtomicReference<>();
3844

3945
public AgentTaskItemListBuilder() {
4046
super();
@@ -58,7 +64,11 @@ public AgentTaskItemListBuilder agent(String name, Object agent) {
5864
exec ->
5965
this.delegate.function(
6066
name,
61-
fn -> fn.function(AgentAdapters.toFunction(exec), DefaultAgenticScope.class)));
67+
fn ->
68+
fn.function(
69+
AgentAdapters.toFunction(
70+
exec, beforeAgentInvocation, afterAgentInvocation),
71+
DefaultAgenticScope.class)));
6272
return self();
6373
}
6474

@@ -96,7 +106,10 @@ public AgentTaskItemListBuilder parallel(String name, Object... agents) {
96106
ex.agentInvoker().name() != null
97107
? ex.agentInvoker().name()
98108
: "branch-" + i + "-" + name;
99-
fork.branch(agentName, AgentAdapters.toFunction(ex), DefaultAgenticScope.class);
109+
fork.branch(
110+
agentName,
111+
AgentAdapters.toFunction(ex, beforeAgentInvocation, afterAgentInvocation),
112+
DefaultAgenticScope.class);
100113
}
101114
});
102115
return self();
@@ -154,4 +167,14 @@ public AgentTaskItemListBuilder switchCase(
154167
this.delegate.switchCase(name, itemsConfigurer);
155168
return self();
156169
}
170+
171+
public AgentTaskItemListBuilder inputFrom(Consumer<AgenticScope> beforeAgentInvocation) {
172+
this.beforeAgentInvocation.set(beforeAgentInvocation);
173+
return self();
174+
}
175+
176+
public AgentTaskItemListBuilder outputAs(Consumer<AgenticScope> afterAgentInvocationConsumer) {
177+
this.afterAgentInvocation.set(afterAgentInvocationConsumer);
178+
return self();
179+
}
157180
}

‎experimental/fluent/agentic/src/main/java/io/serverlessworkflow/fluent/agentic/LoopAgentsBuilder.java‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,9 @@
2222
import io.serverlessworkflow.fluent.func.FuncTaskItemListBuilder;
2323
import java.util.List;
2424
import java.util.UUID;
25+
import java.util.concurrent.atomic.AtomicReference;
2526
import java.util.function.BiPredicate;
27+
import java.util.function.Consumer;
2628
import java.util.function.ObjIntConsumer;
2729
import java.util.function.Predicate;
2830
import java.util.stream.IntStream;
@@ -31,6 +33,10 @@ public class LoopAgentsBuilder {
3133

3234
private final FuncTaskItemListBuilder funcDelegate;
3335
private final ForTaskFunction forTask;
36+
private final AtomicReference<Consumer<AgenticScope>> beforeAgentInvocation =
37+
new AtomicReference<>();
38+
private final AtomicReference<Consumer<AgenticScope>> afterAgentInvocation =
39+
new AtomicReference<>();
3440

3541
private int maxIterations = 1024;
3642

@@ -50,7 +56,11 @@ public LoopAgentsBuilder subAgents(String baseName, Object... agents) {
5056
execs,
5157
(exec, idx) ->
5258
funcDelegate.function(
53-
baseName + "-" + idx, fn -> fn.function(AgentAdapters.toFunction(exec))));
59+
baseName + "-" + idx,
60+
fn ->
61+
fn.function(
62+
AgentAdapters.toFunction(
63+
exec, beforeAgentInvocation, afterAgentInvocation))));
5464
return this;
5565
}
5666

@@ -73,6 +83,16 @@ public LoopAgentsBuilder exitCondition(BiPredicate<AgenticScope, Integer> exitCo
7383
return this;
7484
}
7585

86+
public LoopAgentsBuilder inputFrom(Consumer<AgenticScope> beforeAgentInvocationConsumer) {
87+
this.beforeAgentInvocation.set(beforeAgentInvocationConsumer);
88+
return this;
89+
}
90+
91+
public LoopAgentsBuilder outputAs(Consumer<AgenticScope> afterAgentInvocationConsumer) {
92+
this.afterAgentInvocation.set(afterAgentInvocationConsumer);
93+
return this;
94+
}
95+
7696
public ForTaskFunction build() {
7797
this.forTask.setDo(this.funcDelegate.build());
7898
this.forTask.withCollection(ignored -> IntStream.range(0, maxIterations).boxed().toList());

‎experimental/fluent/agentic/src/test/java/io/serverlessworkflow/fluent/agentic/AgenticBeforeAfterInvokeTest.java‎

Lines changed: 202 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,202 @@
1+
/*
2+
* Copyright 2020-Present The Serverless Workflow Specification Authors
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package io.serverlessworkflow.fluent.agentic;
17+
18+
import static org.junit.jupiter.api.Assertions.assertEquals;
19+
import static org.junit.jupiter.api.Assertions.assertTrue;
20+
import static org.mockito.ArgumentMatchers.eq;
21+
import static org.mockito.Mockito.mock;
22+
import static org.mockito.Mockito.when;
23+
24+
import io.serverlessworkflow.api.types.Workflow;
25+
import io.serverlessworkflow.impl.WorkflowApplication;
26+
import java.util.HashMap;
27+
import java.util.List;
28+
import java.util.Map;
29+
import java.util.concurrent.ExecutionException;
30+
import java.util.concurrent.atomic.AtomicInteger;
31+
import org.junit.jupiter.api.Test;
32+
33+
public class AgenticBeforeAfterInvokeTest {
34+
35+
@Test
36+
void testAgentBeforeAfter() throws ExecutionException, InterruptedException {
37+
Agents.ChatBot chatBot = mock(Agents.ChatBot.class);
38+
when(chatBot.chat(eq("A Great Story"))).thenReturn("Once upon a time...");
39+
40+
Workflow wf =
41+
AgentWorkflowBuilder.workflow()
42+
.tasks(
43+
t ->
44+
t.agent(chatBot)
45+
.inputFrom(
46+
agenticScope -> {
47+
assertTrue(agenticScope.state().containsKey("userInput"));
48+
assertEquals("A Great Story", agenticScope.readState("userInput"));
49+
agenticScope.writeState("testData", "someValue");
50+
})
51+
.outputAs(
52+
agenticScope -> {
53+
assertTrue(agenticScope.state().containsKey("userInput"));
54+
assertEquals("A Great Story", agenticScope.readState("userInput"));
55+
assertTrue(agenticScope.state().containsKey("testData"));
56+
assertEquals("someValue", agenticScope.readState("testData"));
57+
assertEquals("Once upon a time...", agenticScope.readState("input"));
58+
agenticScope.writeState("outputData", "outputValue");
59+
}))
60+
.build();
61+
62+
Map<String, Object> in = new HashMap<>();
63+
in.put("userInput", "A Great Story");
64+
65+
Map<String, Object> out;
66+
try (WorkflowApplication app = WorkflowApplication.builder().build()) {
67+
out = app.workflowDefinition(wf).instance(in).start().get().asMap().orElseThrow();
68+
}
69+
70+
assertTrue(out.containsKey("outputData"));
71+
assertEquals("outputValue", out.get("outputData"));
72+
assertTrue(out.containsKey("userInput"));
73+
assertEquals("A Great Story", out.get("userInput"));
74+
assertTrue(out.containsKey("testData"));
75+
assertEquals("someValue", out.get("testData"));
76+
assertTrue(out.containsKey("input"));
77+
assertEquals("Once upon a time...", out.get("input"));
78+
}
79+
80+
@Test
81+
void testSeqBeforeAfter() throws ExecutionException, InterruptedException {
82+
Agents.CreativeWriter creativeWriter = mock(Agents.CreativeWriter.class);
83+
when(creativeWriter.generateStory((eq("dragons and wizards"))))
84+
.thenReturn("A fantasy story...");
85+
86+
Agents.StyleEditor styleEditor = mock(Agents.StyleEditor.class);
87+
when(styleEditor.editStory(eq("A fantasy story..."), eq("young adults")))
88+
.thenReturn("An engaging fantasy story for young adults...");
89+
90+
Agents.StyleScorer styleScorer = mock(Agents.StyleScorer.class);
91+
when(styleScorer.scoreStyle(eq("An engaging fantasy story for young adults..."), eq("funny")))
92+
.thenReturn(8.0);
93+
94+
Workflow wf =
95+
AgentWorkflowBuilder.workflow()
96+
.tasks(
97+
t ->
98+
t.sequence(creativeWriter, styleEditor, styleScorer)
99+
.inputFrom(
100+
agenticScope -> {
101+
if (agenticScope.readState("step") == null) {
102+
agenticScope.writeState("topic", "dragons and wizards");
103+
agenticScope.writeState("step", "out_1");
104+
} else if (agenticScope.readState("step").equals("in_2")) {
105+
agenticScope.writeState("story", agenticScope.readState("input"));
106+
agenticScope.writeState("style", "young adults");
107+
agenticScope.writeState("step", "out_2");
108+
} else if (agenticScope.readState("step").equals("in_3")) {
109+
agenticScope.writeState("story", agenticScope.readState("input"));
110+
agenticScope.writeState("style", "funny");
111+
agenticScope.writeState("step", "out_3");
112+
} else {
113+
throw new RuntimeException("We should not reach here");
114+
}
115+
})
116+
.outputAs(
117+
agenticScope -> {
118+
if (agenticScope.readState("step").equals("out_1")) {
119+
assertEquals("A fantasy story...", agenticScope.readState("input"));
120+
agenticScope.writeState("step", "in_2");
121+
} else if (agenticScope.readState("step").equals("out_2")) {
122+
assertEquals(
123+
"An engaging fantasy story for young adults...",
124+
agenticScope.readState("input"));
125+
agenticScope.writeState("step", "in_3");
126+
} else if (agenticScope.readState("step").equals("out_3")) {
127+
assertEquals(8.0, agenticScope.readState("input"));
128+
} else {
129+
throw new RuntimeException("We should not reach here");
130+
}
131+
}))
132+
.build();
133+
134+
Map<String, Object> out;
135+
try (WorkflowApplication app = WorkflowApplication.builder().build()) {
136+
out = app.workflowDefinition(wf).instance(Map.of()).start().get().asMap().orElseThrow();
137+
}
138+
139+
assertEquals(8.0, out.get("input"));
140+
assertEquals("An engaging fantasy story for young adults...", out.get("story"));
141+
assertEquals("funny", out.get("style"));
142+
}
143+
144+
@Test
145+
void testParallelBeforeAfter() throws ExecutionException, InterruptedException {
146+
Agents.ChatBot chatBot = mock(Agents.ChatBot.class);
147+
when(chatBot.chat(eq("A Great Story"))).thenReturn("Once upon a time...");
148+
149+
Agents.MovieExpert movieExpert = mock(Agents.MovieExpert.class);
150+
when(movieExpert.findMovie(eq("si-fi"))).thenReturn(List.of("Movie A", "Movie B"));
151+
152+
Workflow wf =
153+
AgentWorkflowBuilder.workflow()
154+
.tasks(
155+
t ->
156+
t.parallel(chatBot, movieExpert)
157+
.inputFrom(
158+
agenticScope -> {
159+
agenticScope.writeState("userInput", "A Great Story");
160+
agenticScope.writeState("mood", "si-fi");
161+
}))
162+
.build();
163+
164+
try (WorkflowApplication app = WorkflowApplication.builder().build()) {
165+
app.workflowDefinition(wf).instance(Map.of()).start().get().asJavaObject();
166+
}
167+
}
168+
169+
@Test
170+
void testLoopBeforeAfter() throws ExecutionException, InterruptedException {
171+
Agents.MovieExpert expert = mock(Agents.MovieExpert.class);
172+
when(expert.findMovie(eq("si-fi"))).thenReturn(List.of("Movie A", "Movie B"));
173+
AtomicInteger atomicInteger = new AtomicInteger(0);
174+
175+
Workflow wf =
176+
AgentWorkflowBuilder.workflow()
177+
.tasks(
178+
d ->
179+
d.loop(
180+
l ->
181+
l.subAgents(expert)
182+
.inputFrom(a -> atomicInteger.incrementAndGet())
183+
.outputAs(
184+
agenticScope ->
185+
agenticScope.writeState("count", atomicInteger.get()))
186+
.exitCondition(exit -> ((Integer) exit.readState("count")) == 10)))
187+
.build();
188+
189+
Map<String, Object> result;
190+
try (WorkflowApplication app = WorkflowApplication.builder().build()) {
191+
result =
192+
app.workflowDefinition(wf)
193+
.instance(Map.of("mood", "si-fi", "count", 0))
194+
.start()
195+
.get()
196+
.asMap()
197+
.orElseThrow();
198+
}
199+
200+
assertEquals(10, result.get("count"));
201+
}
202+
}

0 commit comments

Comments
 (0)