Agents: {{{xxx}}} placeholders to insert multi values in json format; AgentContext: Convenience methods; Agents: Better error handling; Updated prompt strategies; More response handlers;

git-svn-id: svn://home0/lm-lib@3653 71dcc7fd-b1f2-4e5d-bb8e-274146e81173

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

@@ -0,0 +1,39 @@
+package de.extio.lmlib.agent;
+
+import java.util.ArrayList;
+import java.util.function.Function;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import de.extio.lmlib.client.Completion;
+
+public class AccumulateTextAgentResponseHandler implements AgentResponseHandler {
+	
+	private static final Logger LOGGER = LoggerFactory.getLogger(AccumulateTextAgentResponseHandler.class);
+	
+	private final String key;
+	
+	private final Function<String, String> transformer;
+	
+	public AccumulateTextAgentResponseHandler(final String key) {
+		this(key, null);
+	}
+	
+	public AccumulateTextAgentResponseHandler(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);
+		}
+		((ArrayList<Object>) split.context().getContext().computeIfAbsent(this.key, k -> new ArrayList<>())).add(response);
+		return true;
+	}
+	
+}

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

@@ -14,6 +14,9 @@ import java.util.stream.Collectors;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+
 import de.extio.lmlib.client.Client;
 import de.extio.lmlib.client.Completion;
 import de.extio.lmlib.client.Conversation;
@@ -24,6 +27,8 @@ public record Agent(String name, AgentType agentType, ModelCategory modelCategor
 	
 	private static final Logger LOGGER = LoggerFactory.getLogger(Agent.class);
 	
+	private static final ObjectMapper objectMapper = new ObjectMapper();
+	
 	public List<AgentContext> execute(final Client client, final ExecutorService agentExecutorService, final AgentContext context) {
 		final var finishedSplits = Collections.synchronizedList(new ArrayList<Split>());
 		
@@ -56,7 +61,7 @@ public record Agent(String name, AgentType agentType, ModelCategory modelCategor
 							
 							try {
 								if (!this.responseHandler.handle(split, completion)) {
-									context.getGraph().add("⚠");
+									split.context().getGraph().add("⚠");
 									LOGGER.warn("{} Cannot parse response: {}", this.name(), completion.response());
 									continue;
 								}
@@ -65,14 +70,15 @@ public record Agent(String name, AgentType agentType, ModelCategory modelCategor
 								break;
 							}
 							catch (final Exception e) {
-								context.getGraph().add("⚠");
+								split.context().getGraph().add("⚠");
 								LOGGER.warn("{} Cannot parse response: {}", this.name(), completion.response(), e);
 								continue;
 							}
 						}
 						if (!parseable) {
 							LOGGER.warn("{} Response is still not parseable after last attempt", this.name());
-							context.getGraph().add("☢");
+							split.context().getGraph().add("☢");
+							split.context().setError(true);
 							return;
 						}
 					}
@@ -85,6 +91,7 @@ public record Agent(String name, AgentType agentType, ModelCategory modelCategor
 				}
 				catch (final Exception ex) {
 					LOGGER.error("Error during agent execution", ex);
+					split.context().setError(true);
 				}
 			}, agentExecutorService));
 		}
@@ -93,19 +100,24 @@ public record Agent(String name, AgentType agentType, ModelCategory modelCategor
 		
 		var retContexts = finishedSplits
 				.stream()
+				.filter(s -> !s.context().isError())
 				.sorted((s0, s1) -> Integer.compare(s0.index(), s1.index()))
 				.map(Split::context)
-				.toList();
+				.collect(Collectors.toCollection(ArrayList::new));
 		
-		if (this.merger != null) {
-			final int size = retContexts.size();
-			retContexts = this.merger.apply(retContexts);
-			for (final var retContext : retContexts) {
-				retContext.getGraph().add(size + "↣" + retContexts.size());
+		if (!retContexts.isEmpty()) {
+			if (this.merger != null) {
+				final int size = retContexts.size();
+				retContexts = new ArrayList<>(this.merger.apply(retContexts));
+				for (final var retContext : retContexts) {
+					retContext.getGraph().add(size + "↣" + retContexts.size());
+				}
 			}
+			
+			retContexts.forEach(this::nextAgent);
 		}
 		
-		retContexts.forEach(this::chooseNextAgent);
+		retContexts.addAll(finishedSplits.stream().filter(s -> s.context().isError()).map(s -> s.context()).toList());
 		
 		return retContexts;
 	}
@@ -134,7 +146,17 @@ public record Agent(String name, AgentType agentType, ModelCategory modelCategor
 	private List<Split> applyTemplate(final AgentContext context) {
 		String text = this.textTemplate != null ? this.textTemplate : "";
 		for (final var entry : context.getContext().entrySet()) {
-			final var k = "{{" + entry.getKey() + "}}";
+			var k = "{{{" + entry.getKey() + "}}}";
+			if (text.contains(k)) {
+				try {
+					text = text.replace(k, objectMapper.writeValueAsString(entry.getValue()));
+				}
+				catch (final JsonProcessingException e) {
+					LOGGER.warn("Cannot write json", e);
+				}
+			}
+			
+			k = "{{" + entry.getKey() + "}}";
 			if (text.contains(k)) {
 				text = text.replace(k, entry.getValue().stream().map(Object::toString).collect(Collectors.joining("\n\n")));
 			}
@@ -204,7 +226,7 @@ public record Agent(String name, AgentType agentType, ModelCategory modelCategor
 		return conversation;
 	}
 	
-	private void chooseNextAgent(final AgentContext context) {
+	private void nextAgent(final AgentContext context) {
 		context.setNextAgent(this.chooseNext.apply(context));
 		
 		if (context.getNextAgent() != null) {

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

@@ -1,10 +1,13 @@
 package de.extio.lmlib.agent;
 
 import java.util.ArrayList;
+import java.util.Collection;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.stream.Collectors;
 
 import de.extio.lmlib.client.Completion;
 import de.extio.lmlib.client.Conversation;
@@ -25,6 +28,8 @@ public final class AgentContext {
 	
 	private volatile Completion lastCompletion;
 	
+	private volatile boolean error;
+	
 	public AgentContext(final Map<String, Agent> agents) {
 		this.agents = agents;
 		this.context = new ConcurrentHashMap<>();
@@ -42,11 +47,36 @@ public final class AgentContext {
 		this.nextAgent = other.nextAgent;
 		this.requestStatistic = other.requestStatistic;
 		this.lastCompletion = other.lastCompletion;
+		this.error = other.error;
 		synchronized (other.graph) {
 			this.graph.addAll(other.graph);
 		}
 	}
 	
+	public String getStringValue(final String key) {
+		final var values = this.context.get(key);
+		if (values == null || values.isEmpty()) {
+			return null;
+		}
+		return values.getFirst().toString();
+	}
+	
+	public void setStringValue(final String key, final String value) {
+		this.context.put(key, List.of(value));
+	}
+	
+	public List<String> getStringValues(final String key) {
+		final var values = this.context.get(key);
+		if (values == null) {
+			return null;
+		}
+		return values.stream().map(Objects::toString).collect(Collectors.toCollection(ArrayList::new));
+	}
+	
+	public void setStringValues(final String key, final Collection<String> value) {
+		this.context.put(key, List.copyOf(value));
+	}
+	
 	public Map<String, List<? extends Object>> getContext() {
 		return this.context;
 	}
@@ -87,4 +117,12 @@ public final class AgentContext {
 		this.lastCompletion = lastCompletion;
 	}
 	
+	public boolean isError() {
+		return this.error;
+	}
+	
+	public void setError(final boolean error) {
+		this.error = error;
+	}
+	
 }

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

@@ -53,7 +53,7 @@ public class AgentExecutorService implements InitializingBean, DisposableBean {
 		final var branches = agent.execute(client, this.agentExecutorService, context);
 		final var branchTasks = new ArrayList<CompletableFuture<?>>(branches.size());
 		for (final var branchContext : branches) {
-			if (branchContext.getNextAgent() != null) {
+			if (!branchContext.isError() && branchContext.getNextAgent() != null) {
 				final var branchAgent = context.getAgents().get(branchContext.getNextAgent().name());
 				if (branchAgent == null) {
 					throw new IllegalArgumentException("Agent not found: " + branchContext.getNextAgent().name());
@@ -63,7 +63,7 @@ public class AgentExecutorService implements InitializingBean, DisposableBean {
 					try {
 						responses.addAll(this.walkGraph(branchAgent, branchContext));
 					}
-					catch (Exception e) {
+					catch (final Exception e) {
 						LOGGER.error("Error executing agent", e);
 					}
 				}, this.branchExecutorService));

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

@@ -1,6 +1,7 @@
 package de.extio.lmlib.agent;
 
 import java.util.List;
+import java.util.function.Function;
 
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -13,14 +14,25 @@ public class TextAgentResponseHandler implements AgentResponseHandler {
 	
 	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) {
-		split.context().getContext().put(this.key, List.of(completion.response()));
-		LOGGER.debug("Response: {}", completion.response());
+		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/prompt/MistralNemoPromptStrategy.java

@@ -0,0 +1,50 @@
+package de.extio.lmlib.prompt;
+
+import org.springframework.stereotype.Component;
+
+@Component
+public class MistralNemoPromptStrategy implements PromptStrategy {
+	
+	@Override
+	public StringBuilder start(final String instruction, final String question, final String text) {
+		final StringBuilder prompt = new StringBuilder();
+		prompt.append("<s>[INST]");
+		if (! instruction.isEmpty()) {
+			prompt.append(instruction);
+			prompt.append("\n");
+		}
+		if (!question.isEmpty() || !text.isEmpty()) {
+			prompt.append(question);
+			if (!question.isEmpty() && !text.isEmpty()) {
+				prompt.append("\n");
+			}
+			prompt.append(text);
+		}
+		prompt.append("[/INST]");
+		return prompt;
+	}
+	
+	@Override
+	public void continue_(final StringBuilder prompt, final String assistant) {
+		this.next(prompt, assistant, "Continue");
+	}
+	
+	@Override
+	public void next(final StringBuilder prompt, final String assistant, final String user) {
+		prompt.append(assistant);
+		prompt.append("</s>[INST]");
+		prompt.append(user);
+		prompt.append("[/INST]");
+	}
+	
+	@Override
+	public String removeEOT(final String prompt) {
+		return prompt.strip().replace("</s>", "").replace("<|im_end|>", "");
+	}
+	
+	@Override
+	public String getPromptName() {
+		return "mistralNemo";
+	}
+	
+}

src/main/java/de/extio/lmlib/prompt/MistralPromptStrategy.java

@@ -3,12 +3,12 @@ package de.extio.lmlib.prompt;
 import org.springframework.stereotype.Component;
 
 @Component
-public class MixtralPromptStrategy implements PromptStrategy {
+public class MistralPromptStrategy implements PromptStrategy {
 	
 	@Override
 	public StringBuilder start(final String instruction, final String question, final String text) {
 		final StringBuilder prompt = new StringBuilder();
-		prompt.append("<s>[INST] ");
+		prompt.append("<s>[INST]");
 		prompt.append(instruction);
 		if (!instruction.isEmpty() && !question.isEmpty()) {
 			prompt.append("\n");
@@ -29,21 +29,21 @@ public class MixtralPromptStrategy implements PromptStrategy {
 	
 	@Override
 	public void next(final StringBuilder prompt, final String assistant, final String user) {
-		prompt.append("\n\"");
+		prompt.append(" ");
 		prompt.append(assistant);
-		prompt.append("\"</s> [INST] ");
+		prompt.append("</s> [INST] ");
 		prompt.append(user);
 		prompt.append(" [/INST]");
 	}
 	
 	@Override
 	public String removeEOT(final String prompt) {
-		return prompt.strip().replace("</s>", "");
+		return prompt.strip().replace("</s>", "").replace("<|im_end|>", "");
 	}
 	
 	@Override
 	public String getPromptName() {
-		return "mixtral";
+		return "mistral";
 	}
 	
 }