Refactor BaseAgent

README.md

@@ -131,6 +131,10 @@ All language model requests will be cached if you add a CachedClientRepository b
 
 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.
+
 # Usage
 
 1. Setup dependencies, see above

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

@@ -6,6 +6,57 @@ import java.util.function.Function;
 
 import de.extio.lmlib.profile.ModelCategory;
 
-public record Agent(String name, AgentType agentType, ModelCategory modelCategory, String systemPrompt, String textTemplate, AgentResponseHandler responseHandler, Consumer<AgentContext> preProcessor, Consumer<AgentContext> postProcessor, Function<List<AgentContext>, List<AgentContext>> merger, Function<AgentContext, AgentNext> chooseNext) implements BaseAgent {
+public record Agent(String name,
+        AgentType agentType,
+        ModelCategory modelCategory,
+        String systemPrompt,
+        String textTemplate,
+        AgentResponseHandler responseHandler,
+        Consumer<AgentContext> preProcessor,
+        Consumer<AgentContext> postProcessor,
+        Function<List<AgentContext>,
+        List<AgentContext>> merger,
+        Function<AgentContext, AgentNext> chooseNext) implements BaseAgent {
+    
+    public Agent {
+        if (name == null) {
+            throw new IllegalArgumentException("Name cannot be null");
+        }
+        if (agentType == null) {
+            throw new IllegalArgumentException("Agent type cannot be null");
+        }
+    }
+    
+    @Override
+    public void preProcess(final AgentContext context) {
+        if (preProcessor != null) {
+            preProcessor.accept(context);
+        }
+        BaseAgent.super.preProcess(context);
+    }
+    
+    @Override
+    public void postProcess(final AgentContext context) {
+        if (postProcessor != null) {
+            postProcessor.accept(context);
+        }
+        BaseAgent.super.postProcess(context);
+    }
+    
+    @Override
+    public List<AgentContext> merge(final List<AgentContext> contexts) {
+        if (merger != null) {
+            return merger.apply(contexts);
+        }
+        return BaseAgent.super.merge(contexts);
+    }
+    
+    @Override
+    public AgentNext chooseNextAgent(final AgentContext context) {
+        if (chooseNext != null) {
+            return chooseNext.apply(context);
+        }
+        return BaseAgent.super.chooseNextAgent(context);
+    }
     
 }

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

@@ -7,8 +7,6 @@ import java.util.List;
 import java.util.Map.Entry;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
-import java.util.function.Consumer;
-import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import org.slf4j.Logger;
@@ -25,16 +23,16 @@ import de.extio.lmlib.profile.ModelCategory;
 
 public interface BaseAgent {
 	
-	static final Logger LOGGER = LoggerFactory.getLogger(Agent.class);
+	static final Logger LOGGER = LoggerFactory.getLogger(BaseAgent.class);
 	
 	static final ObjectMapper objectMapper = new ObjectMapper();
 	
 	default String name() {
-		throw new IllegalStateException("Name must be implemented in the agent");
+		throw new UnsupportedOperationException("Name must be implemented in the agent");
 	}
 	
 	default AgentType agentType() {
-		throw new IllegalStateException("Agent type must be implemented in the agent");
+		throw new UnsupportedOperationException("Agent type must be implemented in the agent");
 	}
 	
 	default ModelCategory modelCategory() {
@@ -42,34 +40,30 @@ public interface BaseAgent {
 	}
 	
 	default String systemPrompt() {
-		return "";
+		return null;
 	}
 	
 	default String textTemplate() {
-		return "";
+		return null;
 	}
 	
 	default AgentResponseHandler responseHandler() {
-		return (split, completion) -> true;
+		return new TextAgentResponseHandler("response");
 	}
 	
-	default Consumer<AgentContext> preProcessor() {
+	default void preProcess(final AgentContext context) {
+		
+	}
+	
+	default void postProcess(final AgentContext context) {
+		
+	}
+	
+	default List<AgentContext> merge(final List<AgentContext> contexts) {
 		return null;
 	}
 	
-	default Consumer<AgentContext> postProcessor() {
+	default AgentNext chooseNextAgent(final AgentContext context) {
-		return null;
-	}
-	
-	default Function<List<AgentContext>, List<AgentContext>> merger() {
-		return null;
-	}
-	
-	default Function<AgentContext, AgentNext> chooseNext() {
-		return this::chooseNext;
-	}
-	
-	default AgentNext chooseNext(final AgentContext context) {
 		return AgentNext.END;
 	}
 	
@@ -78,9 +72,7 @@ public interface BaseAgent {
 		
 		context.setNextAgent(null);
 		
-		if (this.preProcessor() != null) {
+		this.preProcess(context);
-			this.preProcessor().accept(context);
-		}
 		
 		final var splits = this.applyTemplate(context);
 		final var tasks = new ArrayList<CompletableFuture<?>>(splits.size());
@@ -132,9 +124,7 @@ public interface BaseAgent {
 						}
 					}
 					
-					if (this.postProcessor() != null) {
+					this.postProcess(split.context());
-						this.postProcessor().accept(split.context());
-					}
 					
 					finishedSplits.add(split);
 				}
@@ -155,9 +145,10 @@ public interface BaseAgent {
 				.collect(Collectors.toCollection(ArrayList::new));
 		
 		if (!retContexts.isEmpty()) {
-			if (this.merger() != null) {
+			final var merged = this.merge(List.copyOf(retContexts));
-				final int size = retContexts.size();
+			if (merged != null) {
-				retContexts = new ArrayList<>(this.merger().apply(retContexts));
+				final var size = retContexts.size();
+				retContexts = new ArrayList<>(merged);
 				for (final var retContext : retContexts) {
 					retContext.getGraph().add(size + "↣" + retContexts.size());
 				}
@@ -276,7 +267,7 @@ public interface BaseAgent {
 	}
 	
 	private void nextAgent(final AgentContext context) {
-		context.setNextAgent(this.chooseNext().apply(context));
+		context.setNextAgent(this.chooseNextAgent(context));
 		
 		if (context.getNextAgent() != null) {
 			if (context.getNextAgent().gradingPassed() != null) {

src/test/java/de/extio/lmlib/agent/AgentTest.java

@@ -62,7 +62,7 @@ public class AgentTest {
 						AgentType.CONVERSATION,
 						ModelCategory.MEDIUM,
 						"",
-						"Generate a list of 5 distinct features. Provide no preamble and no explanation. Return the response in JSON format with the following field: { \"features\": [\"Feature 1\", \"Feature 2\", ... ] }",
+						"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())),
@@ -103,7 +103,7 @@ public class AgentTest {
 						context -> 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/Agent.java"))));
+		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);
 		

src/test/resources/application-test.properties

@@ -1,10 +1,10 @@
-profile.model.s=llama3-local
+profile.model.s=local-completion
-profile.model.m=llama3-local
+profile.model.m=local-completion
-profile.model.l=llama3-local
+profile.model.l=local-completion
-profile.model.hot=llama3-local
+profile.model.hot=local-completion
-profile.model.cold=llama3-local
+profile.model.cold=local-completion
 
-agent.threads=4
+agent.threads=2
 
 client.collectStatistics=true
 

src/test/resources/local-completion.properties

@@ -1,8 +1,8 @@
-prompts=llama3
+prompts=chatml
 tokenEncoding=cl100k_base
 maxTokens=2000
 maxContextLength=8000
-temperature=0.2
+temperature=0.4
 topP=1.0
 maxContinuations=1
 modelProvider=OAI_TEXT_COMPLETION