exo/lm-lib
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Implemented agent streaming

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

1
src/main/java/de/extio/lmlib/agent/Agent.java
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,

46
src/main/java/de/extio/lmlib/agent/AgentContext.java
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;
}
}

28
src/main/java/de/extio/lmlib/agent/AgentExecutorService.java
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;
@@ -44,29 +45,42 @@ public class AgentExecutorService implements InitializingBean, DisposableBean {
this.branchExecutorService.close();
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> walk(final BaseAgent agent, final AgentContext context) {
return this.walk(agent, context, null);
}
public List<AgentContext> walkGraph(final BaseAgent agent, final AgentContext context) {
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);

3
src/main/java/de/extio/lmlib/agent/AgentNext.java
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;
}
}

9
src/main/java/de/extio/lmlib/agent/AgentResponseHandler.java
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);
}

23
src/main/java/de/extio/lmlib/agent/BaseAgent.java
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;

39
src/main/java/de/extio/lmlib/agent/TextAgentResponseHandler.java
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;
}
}

8
src/main/java/de/extio/lmlib/agent/AccumulateTextAgentResponseHandler.java → src/main/java/de/extio/lmlib/agent/responsehandler/AccumulateTextAgentResponseHandler.java
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;
}

9
src/main/java/de/extio/lmlib/agent/responsehandler/AgentResponseHandler.java Executable file
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);
}

19
src/main/java/de/extio/lmlib/agent/JsonAgentResponseHandler.java → src/main/java/de/extio/lmlib/agent/responsehandler/JsonAgentResponseHandler.java
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."));
}
}

11
src/main/java/de/extio/lmlib/agent/responsehandler/StreamedAgentResponseHandler.java Normal file
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);
}

65
src/main/java/de/extio/lmlib/agent/responsehandler/TextAgentResponseHandler.java Executable file
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;
}
}