Implemented agent streaming

This commit is contained in:
extio-de
2025-06-28 16:55:26 +02:00
parent 365043b5a0
commit 9ca590a8d4
14 changed files with 259 additions and 87 deletions

View File

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

View File

@@ -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,

View File

@@ -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 volatile boolean skipNextCompletion;
private Consumer<AgentContext> agentContextUpdateConsumer;
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;
}
}

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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

View File

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