From 9ca590a8d45eec0ce74e5993fa7cf1930ad4b442 Mon Sep 17 00:00:00 2001 From: extio-de Date: Sat, 28 Jun 2025 16:55:26 +0200 Subject: [PATCH] Implemented agent streaming --- README.md | 30 +++++++-- src/main/java/de/extio/lmlib/agent/Agent.java | 1 + .../de/extio/lmlib/agent/AgentContext.java | 46 ++++++++++--- .../lmlib/agent/AgentExecutorService.java | 28 ++++++-- .../java/de/extio/lmlib/agent/AgentNext.java | 3 + .../lmlib/agent/AgentResponseHandler.java | 9 --- .../java/de/extio/lmlib/agent/BaseAgent.java | 23 ++++++- .../lmlib/agent/TextAgentResponseHandler.java | 39 ----------- .../AccumulateTextAgentResponseHandler.java | 8 ++- .../responsehandler/AgentResponseHandler.java | 9 +++ .../JsonAgentResponseHandler.java | 19 ++++-- .../StreamedAgentResponseHandler.java | 11 ++++ .../TextAgentResponseHandler.java | 65 +++++++++++++++++++ .../java/de/extio/lmlib/agent/AgentTest.java | 55 ++++++++++++++-- 14 files changed, 259 insertions(+), 87 deletions(-) delete mode 100755 src/main/java/de/extio/lmlib/agent/AgentResponseHandler.java delete mode 100755 src/main/java/de/extio/lmlib/agent/TextAgentResponseHandler.java rename src/main/java/de/extio/lmlib/agent/{ => responsehandler}/AccumulateTextAgentResponseHandler.java (74%) create mode 100755 src/main/java/de/extio/lmlib/agent/responsehandler/AgentResponseHandler.java rename src/main/java/de/extio/lmlib/agent/{ => responsehandler}/JsonAgentResponseHandler.java (79%) create mode 100644 src/main/java/de/extio/lmlib/agent/responsehandler/StreamedAgentResponseHandler.java create mode 100755 src/main/java/de/extio/lmlib/agent/responsehandler/TextAgentResponseHandler.java diff --git a/README.md b/README.md index b38908c..daa0367 100644 --- a/README.md +++ b/README.md @@ -123,7 +123,7 @@ Example profile: src/test/resources/llama3.3-70b-azure.properties Client service: clientService.getClient(ModelCategory.MEDIUM); Client usage details: CompletionClientTest and AzureAiClientTest -Agentic flow: AgentTest#agenticFlow +Agentic flow: AgentTest#agenticFlow and AgentTest#streamedAgenticFlow ## Completions and Streaming @@ -143,10 +143,6 @@ You can intercept every Conversation and Completion by implementing CompletionIn This is for example useful for accounting purposes or to implement a circuit breaker. The client will autowire all implementations and intercept requests accordingly. -# Agentic Flows - -This library contains a graph executor for implementing comprehensive processes using agent-based design patterns. See class AgentTest for examples. - # Grader The library includes a grader component that evaluates the correctness of agent responses. @@ -155,6 +151,30 @@ Example: assertTrue(Grader.assessScoreBinary("Does the following blog post mention the feature " + feature, resultContext.getStringValue("blogPost"), this.clientService)); +# Agentic Flows + +This library contains a graph executor for implementing comprehensive processes using agent-based design patterns. + +You can implement your Agents either by creating an Agent record or by extending BaseAgent - choose the approach that best fits your needs. +The agentic flow is started with an AgentContext and an initial agent. +Example: + + final var context = new AgentContext(List.of(new Agent1(), new Agent2())); + context.getContext().setStringValue("key", "Payload here"); + context.getContext().setStringValues("key2", List.of("All context keys", "support collections", "for aggregation purposes")); + context.getContext().setValue("key3", new RandomDatatype()); + final var resultContexts = agentExecutorService.walk("Agent1", context) + final var resultContext = resultContexts.getFirst(); // You may receive multiple contexts if your graph branches out + LOGGER.info(resultContext.getStringValue("blogPost") + "\n\n" + resultContext.getGraph().toString() + "\n" + resultContext.getRequestStatistic().toString()); + +Streaming is also supported. + +See class AgentTest for examples. + +# Agent implementations + + + # Announcement generated by AgentTest#agenticFlow ## Introducing Advanced Agent Orchestration: Dynamic Conversations & Parallel Processing diff --git a/src/main/java/de/extio/lmlib/agent/Agent.java b/src/main/java/de/extio/lmlib/agent/Agent.java index 7692b0c..08c085a 100755 --- a/src/main/java/de/extio/lmlib/agent/Agent.java +++ b/src/main/java/de/extio/lmlib/agent/Agent.java @@ -4,6 +4,7 @@ import java.util.List; import java.util.function.Consumer; import java.util.function.Function; +import de.extio.lmlib.agent.responsehandler.AgentResponseHandler; import de.extio.lmlib.profile.ModelCategory; public record Agent(String name, diff --git a/src/main/java/de/extio/lmlib/agent/AgentContext.java b/src/main/java/de/extio/lmlib/agent/AgentContext.java index c4982bb..b626255 100755 --- a/src/main/java/de/extio/lmlib/agent/AgentContext.java +++ b/src/main/java/de/extio/lmlib/agent/AgentContext.java @@ -7,6 +7,7 @@ import java.util.List; import java.util.Map; import java.util.Objects; import java.util.concurrent.ConcurrentHashMap; +import java.util.function.Consumer; import java.util.stream.Collectors; import de.extio.lmlib.client.Completion; @@ -14,23 +15,31 @@ import de.extio.lmlib.client.Conversation; public final class AgentContext { - private final Map agents; + // Exposed private final Map> context; + private volatile boolean streaming; + + private volatile boolean error; + + private volatile boolean skipNextCompletion; + + private volatile Conversation conversation; + private final AgentRequestStatistic requestStatistic; private final List graph = Collections.synchronizedList(new ArrayList<>()); - private volatile Conversation conversation; + // Internal + + private final Map agents; private volatile AgentNext nextAgent; private volatile Completion lastCompletion; - private volatile boolean error; - - private volatile boolean skipNextCompletion; + private Consumer agentContextUpdateConsumer; public AgentContext(final Collection agents) { this.agents = agents.stream().collect(Collectors.toMap(BaseAgent::name, agent -> agent)); @@ -55,8 +64,10 @@ public final class AgentContext { this.nextAgent = other.nextAgent; this.requestStatistic = other.requestStatistic; this.lastCompletion = other.lastCompletion; + this.agentContextUpdateConsumer = other.agentContextUpdateConsumer; this.error = other.error; this.skipNextCompletion = other.skipNextCompletion; + this.streaming = other.streaming; synchronized (other.graph) { this.graph.addAll(other.graph); } @@ -111,7 +122,7 @@ public final class AgentContext { return this.context; } - public Map getAgents() { + Map getAgents() { return this.agents; } @@ -127,11 +138,11 @@ public final class AgentContext { this.conversation = conversation; } - public AgentNext getNextAgent() { + AgentNext getNextAgent() { return this.nextAgent; } - public void setNextAgent(final AgentNext next) { + void setNextAgent(final AgentNext next) { this.nextAgent = next; } @@ -139,11 +150,11 @@ public final class AgentContext { return this.graph; } - public Completion getLastCompletion() { + Completion getLastCompletion() { return this.lastCompletion; } - public void setLastCompletion(final Completion lastCompletion) { + void setLastCompletion(final Completion lastCompletion) { this.lastCompletion = lastCompletion; } @@ -163,4 +174,19 @@ public final class AgentContext { this.skipNextCompletion = skipNextCompletion; } + Consumer getAgentContextUpdateConsumer() { + return agentContextUpdateConsumer; + } + + void setAgentContextUpdateConsumer(final Consumer agentContextUpdateConsumer) { + this.agentContextUpdateConsumer = agentContextUpdateConsumer; + } + + public boolean isStreaming() { + return streaming; + } + + public void setStreaming(final boolean streaming) { + this.streaming = streaming; + } } diff --git a/src/main/java/de/extio/lmlib/agent/AgentExecutorService.java b/src/main/java/de/extio/lmlib/agent/AgentExecutorService.java index 7ce2ec0..267a4f0 100755 --- a/src/main/java/de/extio/lmlib/agent/AgentExecutorService.java +++ b/src/main/java/de/extio/lmlib/agent/AgentExecutorService.java @@ -6,6 +6,7 @@ import java.util.List; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; +import java.util.function.Consumer; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -44,29 +45,42 @@ public class AgentExecutorService implements InitializingBean, DisposableBean { this.branchExecutorService.close(); this.agentExecutorService.close(); } - - public List walkGraph(final String firstAgent, final AgentContext context) { + + public List walk(final String firstAgent, final AgentContext context) { + return this.walk(firstAgent, context, null); + } + + public List walk(final String firstAgent, final AgentContext context, final Consumer agentContextUpdateConsumer) { final var agent = context.getAgents().get(firstAgent); if (agent == null) { throw new IllegalArgumentException("Agent not found: " + firstAgent); } - LOGGER.debug("Starting agent: {}", agent.name()); - return this.walkGraph(agent, context); + + return this.walk(agent, context, agentContextUpdateConsumer); + } + + public List walk(final BaseAgent agent, final AgentContext context) { + return this.walk(agent, context, null); } - public List walkGraph(final BaseAgent agent, final AgentContext context) { + public List walk(final BaseAgent agent, final AgentContext context, final Consumer agentContextUpdateConsumer) { LOGGER.debug("Agent: {}", agent.name()); Client client = null; if (agent.agentType() != AgentType.PROCESSING_ONLY) { client = this.clientService.getClient(agent.modelCategory()); } - + context.setAgentContextUpdateConsumer(agentContextUpdateConsumer); + final var branches = agent.execute(client, this.agentExecutorService, context); final var responses = Collections.synchronizedList(new ArrayList()); final var branchTasks = new ArrayList>(branches.size()); for (final var branchContext : branches) { + if (agentContextUpdateConsumer != null) { + agentContextUpdateConsumer.accept(branchContext); + } + if (!branchContext.isError() && branchContext.getNextAgent() != null) { final var branchAgent = context.getAgents().get(branchContext.getNextAgent().name()); if (branchAgent == null) { @@ -75,7 +89,7 @@ public class AgentExecutorService implements InitializingBean, DisposableBean { branchTasks.add(CompletableFuture.runAsync(() -> { try { - responses.addAll(this.walkGraph(branchAgent, branchContext)); + responses.addAll(this.walk(branchAgent, branchContext, agentContextUpdateConsumer)); } catch (final Exception e) { LOGGER.error("Error executing agent", e); diff --git a/src/main/java/de/extio/lmlib/agent/AgentNext.java b/src/main/java/de/extio/lmlib/agent/AgentNext.java index 48f71c4..bb22053 100755 --- a/src/main/java/de/extio/lmlib/agent/AgentNext.java +++ b/src/main/java/de/extio/lmlib/agent/AgentNext.java @@ -4,4 +4,7 @@ public final record AgentNext(String name, Boolean gradingPassed) { public static final AgentNext END = new AgentNext(null, null); + public static AgentNext end(final AgentContext context) { + return END; + } } diff --git a/src/main/java/de/extio/lmlib/agent/AgentResponseHandler.java b/src/main/java/de/extio/lmlib/agent/AgentResponseHandler.java deleted file mode 100755 index f1d20c2..0000000 --- a/src/main/java/de/extio/lmlib/agent/AgentResponseHandler.java +++ /dev/null @@ -1,9 +0,0 @@ -package de.extio.lmlib.agent; - -import de.extio.lmlib.client.Completion; - -public interface AgentResponseHandler { - - boolean handle(Split split, Completion completion); - -} diff --git a/src/main/java/de/extio/lmlib/agent/BaseAgent.java b/src/main/java/de/extio/lmlib/agent/BaseAgent.java index 7caaee4..86e0fd7 100644 --- a/src/main/java/de/extio/lmlib/agent/BaseAgent.java +++ b/src/main/java/de/extio/lmlib/agent/BaseAgent.java @@ -15,6 +15,9 @@ import org.slf4j.LoggerFactory; import com.fasterxml.jackson.core.JsonProcessingException; import com.fasterxml.jackson.databind.ObjectMapper; +import de.extio.lmlib.agent.responsehandler.AgentResponseHandler; +import de.extio.lmlib.agent.responsehandler.StreamedAgentResponseHandler; +import de.extio.lmlib.agent.responsehandler.TextAgentResponseHandler; import de.extio.lmlib.client.Client; import de.extio.lmlib.client.Completion; import de.extio.lmlib.client.Conversation; @@ -93,15 +96,29 @@ public interface BaseAgent { final var conversation = this.setupConversation(split); LOGGER.debug("Conversation: {}", conversation); - Completion completion = null; boolean parseable = false; for (int i = 0; i < 2; i++) { - completion = client.conversation(this.modelCategory(), conversation); + Completion completion = null; + if (split.context().isStreaming()) { + if (this.responseHandler() instanceof final StreamedAgentResponseHandler streamedResponseHandler) { + streamedResponseHandler.beforeStream(split.context()); + } + completion = client.streamConversation(this.modelCategory(), conversation, chunk -> { + if (this.responseHandler() instanceof final StreamedAgentResponseHandler streamedResponseHandler) { + if (streamedResponseHandler.handleChunk(chunk, split.context()) && split.context().getAgentContextUpdateConsumer() != null) { + split.context().getAgentContextUpdateConsumer().accept(split.context()); + } + } + }); + } + else { + completion = client.conversation(this.modelCategory(), conversation); + } split.context().setLastCompletion(completion); split.context().getRequestStatistic().add(completion); try { - if (!this.responseHandler().handle(split, completion)) { + if (!this.responseHandler().handle(completion, split.context())) { split.context().getGraph().add("⚠"); LOGGER.warn("{} Cannot parse response: {}", this.name(), completion.response()); continue; diff --git a/src/main/java/de/extio/lmlib/agent/TextAgentResponseHandler.java b/src/main/java/de/extio/lmlib/agent/TextAgentResponseHandler.java deleted file mode 100755 index 24dd356..0000000 --- a/src/main/java/de/extio/lmlib/agent/TextAgentResponseHandler.java +++ /dev/null @@ -1,39 +0,0 @@ -package de.extio.lmlib.agent; - -import java.util.List; -import java.util.function.Function; - -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; - -import de.extio.lmlib.client.Completion; - -public class TextAgentResponseHandler implements AgentResponseHandler { - - private static final Logger LOGGER = LoggerFactory.getLogger(TextAgentResponseHandler.class); - - private final String key; - - private final Function transformer; - - public TextAgentResponseHandler(final String key) { - this(key, null); - } - - public TextAgentResponseHandler(final String key, final Function transformer) { - this.key = key; - this.transformer = transformer; - } - - @Override - public boolean handle(final Split split, final Completion completion) { - var response = completion.response(); - LOGGER.debug("Response: {}", response); - if (this.transformer != null) { - response = this.transformer.apply(response); - } - split.context().getContext().put(this.key, List.of(response)); - return true; - } - -} diff --git a/src/main/java/de/extio/lmlib/agent/AccumulateTextAgentResponseHandler.java b/src/main/java/de/extio/lmlib/agent/responsehandler/AccumulateTextAgentResponseHandler.java similarity index 74% rename from src/main/java/de/extio/lmlib/agent/AccumulateTextAgentResponseHandler.java rename to src/main/java/de/extio/lmlib/agent/responsehandler/AccumulateTextAgentResponseHandler.java index f9cefb1..d66ddf5 100644 --- a/src/main/java/de/extio/lmlib/agent/AccumulateTextAgentResponseHandler.java +++ b/src/main/java/de/extio/lmlib/agent/responsehandler/AccumulateTextAgentResponseHandler.java @@ -1,4 +1,4 @@ -package de.extio.lmlib.agent; +package de.extio.lmlib.agent.responsehandler; import java.util.ArrayList; import java.util.function.Function; @@ -6,6 +6,7 @@ import java.util.function.Function; import org.slf4j.Logger; import org.slf4j.LoggerFactory; +import de.extio.lmlib.agent.AgentContext; import de.extio.lmlib.client.Completion; public class AccumulateTextAgentResponseHandler implements AgentResponseHandler { @@ -25,14 +26,15 @@ public class AccumulateTextAgentResponseHandler implements AgentResponseHandler this.transformer = transformer; } + @SuppressWarnings("unchecked") @Override - public boolean handle(final Split split, final Completion completion) { + public boolean handle(final Completion completion, final AgentContext context) { var response = completion.response(); LOGGER.debug("Response: {}", response); if (this.transformer != null) { response = this.transformer.apply(response); } - ((ArrayList) split.context().getContext().computeIfAbsent(this.key, k -> new ArrayList<>())).add(response); + ((ArrayList) context.getContext().computeIfAbsent(this.key, k -> new ArrayList<>())).add(response); return true; } diff --git a/src/main/java/de/extio/lmlib/agent/responsehandler/AgentResponseHandler.java b/src/main/java/de/extio/lmlib/agent/responsehandler/AgentResponseHandler.java new file mode 100755 index 0000000..c6e37af --- /dev/null +++ b/src/main/java/de/extio/lmlib/agent/responsehandler/AgentResponseHandler.java @@ -0,0 +1,9 @@ +package de.extio.lmlib.agent.responsehandler; + +import de.extio.lmlib.agent.AgentContext; +import de.extio.lmlib.client.Completion; + +public interface AgentResponseHandler { + + boolean handle(Completion completion, AgentContext context); +} diff --git a/src/main/java/de/extio/lmlib/agent/JsonAgentResponseHandler.java b/src/main/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandler.java similarity index 79% rename from src/main/java/de/extio/lmlib/agent/JsonAgentResponseHandler.java rename to src/main/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandler.java index 2a44f48..6c9fb9e 100755 --- a/src/main/java/de/extio/lmlib/agent/JsonAgentResponseHandler.java +++ b/src/main/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandler.java @@ -1,4 +1,4 @@ -package de.extio.lmlib.agent; +package de.extio.lmlib.agent.responsehandler; import java.util.Iterator; import java.util.LinkedHashSet; @@ -13,7 +13,9 @@ import com.fasterxml.jackson.databind.DeserializationFeature; import com.fasterxml.jackson.databind.JsonNode; import com.fasterxml.jackson.databind.ObjectMapper; +import de.extio.lmlib.agent.AgentContext; import de.extio.lmlib.client.Completion; +import de.extio.lmlib.client.Conversation; public class JsonAgentResponseHandler implements AgentResponseHandler { @@ -43,7 +45,7 @@ public class JsonAgentResponseHandler implements AgentResponseHandler { } @Override - public boolean handle(final Split split, final Completion completion) { + public boolean handle(final Completion completion, final AgentContext context) { final int firstCurly = completion.response().indexOf("{"); final int lastCurly = completion.response().lastIndexOf("}"); final var response = completion.response().substring(firstCurly, lastCurly + 1); @@ -53,6 +55,7 @@ public class JsonAgentResponseHandler implements AgentResponseHandler { } catch (final Exception e) { LOGGER.warn("Cannot parse json response: {}", completion.response()); + addJsonResponseErrorPrompt(context); return false; } @@ -63,18 +66,26 @@ public class JsonAgentResponseHandler implements AgentResponseHandler { if (entry.getValue().isArray()) { final var items = new LinkedHashSet<>(); entry.getValue().elements().forEachRemaining(item -> items.add(item.asText())); - split.context().getContext().put(entry.getKey(), List.copyOf(items)); + context.getContext().put(entry.getKey(), List.copyOf(items)); hasResponse = true; LOGGER.debug("Response: {} {}", entry.getKey(), items); } else if (entry.getValue().isTextual() || entry.getValue().isNumber() || entry.getValue().isBoolean()) { - split.context().getContext().put(entry.getKey(), List.of(entry.getValue().asText())); + context.setStringValue(entry.getKey(), entry.getValue().asText()); hasResponse = true; LOGGER.debug("Response: {} {}", entry.getKey(), entry.getValue().asText()); } } + if (!hasResponse) { + addJsonResponseErrorPrompt(context); + } return hasResponse; } + + private void addJsonResponseErrorPrompt(final AgentContext context) { + final var turn = context.getConversation().getConversation().getLast(); + context.getConversation().replaceTurn(new Conversation.Turn(turn.type(), turn.text() + "\n\nThe previous response could not be fully processed or validated. Please make sure to format the response in valid JSON syntax with properly escaped characters.")); + } } diff --git a/src/main/java/de/extio/lmlib/agent/responsehandler/StreamedAgentResponseHandler.java b/src/main/java/de/extio/lmlib/agent/responsehandler/StreamedAgentResponseHandler.java new file mode 100644 index 0000000..7273697 --- /dev/null +++ b/src/main/java/de/extio/lmlib/agent/responsehandler/StreamedAgentResponseHandler.java @@ -0,0 +1,11 @@ +package de.extio.lmlib.agent.responsehandler; + +import de.extio.lmlib.agent.AgentContext; + +public interface StreamedAgentResponseHandler extends AgentResponseHandler { + + void beforeStream(AgentContext context); + + boolean handleChunk(String chunk, AgentContext context); + +} diff --git a/src/main/java/de/extio/lmlib/agent/responsehandler/TextAgentResponseHandler.java b/src/main/java/de/extio/lmlib/agent/responsehandler/TextAgentResponseHandler.java new file mode 100755 index 0000000..10cb722 --- /dev/null +++ b/src/main/java/de/extio/lmlib/agent/responsehandler/TextAgentResponseHandler.java @@ -0,0 +1,65 @@ +package de.extio.lmlib.agent.responsehandler; + +import java.util.Objects; +import java.util.function.Function; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import de.extio.lmlib.agent.AgentContext; +import de.extio.lmlib.client.Completion; + +public class TextAgentResponseHandler implements StreamedAgentResponseHandler { + + public static final String UPDATE_KEY = "chunkUpdateKey"; + + private static final Logger LOGGER = LoggerFactory.getLogger(TextAgentResponseHandler.class); + + private final String key; + + private final String chunkKey; + + private final Function transformer; + + private final Function chunkTransformer; + + public TextAgentResponseHandler(final String key) { + this(key, null, null); + } + + public TextAgentResponseHandler(final String key, final Function transformer, final Function chunkTransformer) { + this.key = key; + this.chunkKey = key.concat("_chunk"); + this.transformer = transformer; + this.chunkTransformer = chunkTransformer; + } + + @Override + public boolean handle(final Completion completion, final AgentContext context) { + var response = completion.response(); + LOGGER.debug("Response: {}", response); + if (this.transformer != null) { + response = this.transformer.apply(response); + } + context.setStringValue(this.key, response); + return true; + } + + @Override + public void beforeStream(final AgentContext context) { + context.setStringValue(this.key, ""); + } + + @Override + public boolean handleChunk(final String chunk_, final AgentContext context) { + var chunk = chunk_; + if (this.chunkTransformer != null) { + chunk = this.chunkTransformer.apply(chunk); + } + context.setStringValue(this.chunkKey, chunk); + context.setStringValue(this.key, Objects.requireNonNullElse(context.getStringValue(this.key), "".concat(chunk))); + context.setStringValue(UPDATE_KEY, this.key); + return true; + } + +} diff --git a/src/test/java/de/extio/lmlib/agent/AgentTest.java b/src/test/java/de/extio/lmlib/agent/AgentTest.java index 0afcafa..f3ca591 100644 --- a/src/test/java/de/extio/lmlib/agent/AgentTest.java +++ b/src/test/java/de/extio/lmlib/agent/AgentTest.java @@ -8,7 +8,9 @@ import java.nio.file.Files; import java.nio.file.Path; import java.util.List; import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Disabled; import org.junit.jupiter.api.Test; import org.slf4j.Logger; @@ -20,10 +22,11 @@ import org.springframework.boot.test.context.SpringBootTest.WebEnvironment; import org.springframework.context.annotation.ComponentScan; import org.springframework.context.annotation.PropertySource; +import de.extio.lmlib.agent.responsehandler.JsonAgentResponseHandler; +import de.extio.lmlib.agent.responsehandler.TextAgentResponseHandler; import de.extio.lmlib.client.ClientService; import de.extio.lmlib.grader.Grader; import de.extio.lmlib.profile.ModelCategory; - import io.netty.util.internal.ThreadLocalRandom; /** @@ -67,7 +70,7 @@ public class AgentTest { "Generate a list of 5 distinct functional features. In scope is business logic, not code mechanics. Provide no preamble and no explanation. Return the response in JSON format with the following field: { \"features\": [\"Feature 1\", \"Feature 2\", ... ] }", new JsonAgentResponseHandler(), null, - context -> context.getContext().get("features").forEach(feature -> LOGGER.info(feature.toString())), + context -> context.getStringValues("features").forEach(feature -> LOGGER.info(feature)), null, context -> new AgentNext("FeatureAnalyzer", null)), @@ -98,16 +101,16 @@ public class AgentTest { {{featureDescription}}""", new TextAgentResponseHandler("blogPost"), context -> { - context.getContext().put("topics", List.of(String.join(", ", context.getContext().get("features").stream().map(String::valueOf).toArray(size -> new String[size])))); + context.setStringValues("topics", List.of(String.join(", ", context.getStringValues("features").stream().toArray(size -> new String[size])))); }, null, null, - context -> AgentNext.END)); + AgentNext::end)); final var context = new AgentContext(agents); context.getContext().put("code", List.of(Files.readString(Path.of("src/main/java/de/extio/lmlib/agent/BaseAgent.java")))); - final var resultContexts = this.agentExecutor.walkGraph(agents.get("CodeSummarizer"), context); + final var resultContexts = this.agentExecutor.walk(agents.get("CodeSummarizer"), context); final var resultContext = resultContexts.getFirst(); LOGGER.info(resultContext.getContext().get("blogPost").toString() + "\n\n" + resultContext.getGraph().toString() + "\n" + resultContext.getRequestStatistic().toString()); @@ -117,9 +120,47 @@ public class AgentTest { assertEquals(1, resultContext.getContext().get("blogPost").size()); assertEquals(8, resultContext.getRequestStatistic().getRequests().get()); for (final var feature : resultContext.getContext().get("features")) { - assertTrue(Grader.assessScoreBinary("Does the following blog post mention the feature " + feature, resultContext.getContext().get("blogPost").getFirst().toString(), this.clientService)); + assertTrue(Grader.assessScoreBinary("Does the following blog post mention the feature " + feature, resultContext.getStringValue("blogPost"), this.clientService)); } } + + @Test + void streamedAgenticFlow() throws IOException { + final var agents = Map.of( + "CodeSummarizer", + new Agent("CodeSummarizer", + AgentType.START_CONVERSATION, + ModelCategory.MEDIUM, + "Generate a summary of the following Java source code.", + "The source code is:\n{{code}}", + new TextAgentResponseHandler("summary"), + null, + context -> { + context.setStringValue("summary", context.getStringValue("summary") + "\npostprocessed"); + }, + null, + AgentNext::end)); + + final var context = new AgentContext(agents); + context.setStreaming(true); + context.getContext().put("code", List.of(Files.readString(Path.of("src/main/java/de/extio/lmlib/agent/BaseAgent.java")))); + + final var lastStreamedSummary = new AtomicReference(); + final var resultContexts = this.agentExecutor.walk(agents.get("CodeSummarizer"), context, contextUpdate -> { + if ("summary".equals(contextUpdate.getStringValue(TextAgentResponseHandler.UPDATE_KEY))) { // UPDATE_KEY helps you to identify the streamed chunk + System.out.print(contextUpdate.getStringValue("summary_chunk")); // Single chunks can be accessed via _chunk suffix, depends on the StreamedAgentResponseHandler implementation + lastStreamedSummary.set(contextUpdate.getStringValue("summary")); // The current key value can also be accessed (e.g. concatenated, depends on the StreamedAgentResponseHandler implementation) + } + }); + + assertEquals(1, resultContexts.size()); + assertEquals(lastStreamedSummary.get(), resultContexts.getFirst().getStringValue("summary")); + assertTrue(resultContexts.getFirst().getStringValue("summary").endsWith("postprocessed")); + + System.out.println(); + LOGGER.info("### final result"); + LOGGER.info(resultContexts.getFirst().toString()); + } @Test void doubleBranch() { @@ -187,7 +228,7 @@ public class AgentTest { final var context = new AgentContext(agents); - final var resultContexts = this.agentExecutor.walkGraph(agents.get("FirstLevel"), context); + final var resultContexts = this.agentExecutor.walk(agents.get("FirstLevel"), context); assertEquals(25, resultContexts.size()); assertTrue(resultContexts.getFirst().getRequestStatistic().getRequests().get() >= 56);