Implemented agent streaming

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

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,

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<String, ? extends BaseAgent> agents;
+	// Exposed
 	
 	private final Map<String, List<? extends Object>> context;
 	
+	private volatile boolean streaming;
+	
+	private volatile boolean error;
+	
+	private volatile boolean skipNextCompletion;
+	
+	private volatile Conversation conversation;
+	
 	private final AgentRequestStatistic requestStatistic;
 	
 	private final List<String> graph = Collections.synchronizedList(new ArrayList<>());
 	
-	private volatile Conversation conversation;
+	// Internal
+	
+	private final Map<String, ? extends BaseAgent> agents;
 	
 	private volatile AgentNext nextAgent;
 	
 	private volatile Completion lastCompletion;
 	
-	private volatile boolean error;
+	private Consumer<AgentContext> agentContextUpdateConsumer;
-	
-	private volatile boolean skipNextCompletion;
 	
 	public AgentContext(final Collection<? extends BaseAgent> 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<String, ? extends BaseAgent> getAgents() {
+	Map<String, ? extends BaseAgent> 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<AgentContext> getAgentContextUpdateConsumer() {
+		return agentContextUpdateConsumer;
+	}
+	
+	void setAgentContextUpdateConsumer(final Consumer<AgentContext> agentContextUpdateConsumer) {
+		this.agentContextUpdateConsumer = agentContextUpdateConsumer;
+	}
+	
+	public boolean isStreaming() {
+		return streaming;
+	}
+	
+	public void setStreaming(final boolean streaming) {
+		this.streaming = streaming;
+	}
 }

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;
@@ -45,28 +46,41 @@ public class AgentExecutorService implements InitializingBean, DisposableBean {
 		this.agentExecutorService.close();
 	}
 
-	public List<AgentContext> walkGraph(final String firstAgent, final AgentContext context) {
+	public List<AgentContext> walk(final String firstAgent, final AgentContext context) {
+		return this.walk(firstAgent, context, null);
+	}
+
+	public List<AgentContext> walk(final String firstAgent, final AgentContext context, final Consumer<AgentContext> 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<AgentContext> walkGraph(final BaseAgent agent, final AgentContext context) {
+	public List<AgentContext> walk(final BaseAgent agent, final AgentContext context) {
+		return this.walk(agent, context, null);
+	}
+	
+	public List<AgentContext> walk(final BaseAgent agent, final AgentContext context, final Consumer<AgentContext> 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<AgentContext>());
 		final var branchTasks = new ArrayList<CompletableFuture<?>>(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);

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;
+    }
 }

src/main/java/de/extio/lmlib/agent/AgentResponseHandler.java

@@ -1,9 +0,0 @@
-package de.extio.lmlib.agent;
-
-import de.extio.lmlib.client.Completion;
-
-public interface AgentResponseHandler {
-	
-	boolean handle(Split split, Completion completion);
-	
-}

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 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;

src/main/java/de/extio/lmlib/agent/TextAgentResponseHandler.java

@@ -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<String, String> transformer;
-	
-	public TextAgentResponseHandler(final String key) {
-		this(key, null);
-	}
-	
-	public TextAgentResponseHandler(final String key, final Function<String, String> 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;
-	}
-	
-}

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<Object>) split.context().getContext().computeIfAbsent(this.key, k -> new ArrayList<>())).add(response);
+		((ArrayList<Object>) context.getContext().computeIfAbsent(this.key, k -> new ArrayList<>())).add(response);
 		return true;
 	}
 	

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);
+}

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."));
+	}
+	
 }

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);
+    
+}

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<String, String> transformer;
+	
+	private final Function<String, String> chunkTransformer;
+	
+	public TextAgentResponseHandler(final String key) {
+		this(key, null, null);
+	}
+	
+	public TextAgentResponseHandler(final String key, final Function<String, String> transformer, final Function<String, String> 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;
+	}
+	
+}

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,10 +120,48 @@ 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<String>();
+		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() {
 		final var agents = Map.of(
@@ -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);