lm-lib first commit;

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

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

@@ -0,0 +1,238 @@
+package de.extio.lmlib.agent;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.HashMap;
+import java.util.List;
+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;
+import org.slf4j.LoggerFactory;
+
+import de.extio.lmlib.client.Client;
+import de.extio.lmlib.client.Completion;
+import de.extio.lmlib.client.Conversation;
+import de.extio.lmlib.client.Conversation.TurnType;
+import de.extio.lmlib.client.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) {
+	
+	private static final Logger LOGGER = LoggerFactory.getLogger(Agent.class);
+	
+	public List<AgentContext> execute(final Client client, final ExecutorService agentExecutorService, final AgentContext context) {
+		final var retContexts = Collections.synchronizedList(new ArrayList<AgentContext>());
+		
+		context.setNextAgent(null);
+		
+		if (this.preProcessor != null) {
+			this.preProcessor.accept(context);
+		}
+		
+		final var splits = this.applyTemplate(context);
+		final var tasks = new ArrayList<CompletableFuture<?>>(splits.size());
+		for (final var split : splits) {
+			tasks.add(CompletableFuture.runAsync(() -> {
+				split.context().getGraph().add(this.modelCategory.getShortName());
+				split.context().getGraph().add(this.name());
+				
+				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);
+					split.context().getRequestStatistic().add(completion);
+					
+					try {
+						if (!this.responseHandler.handle(split, completion)) {
+							context.getGraph().add("⚠");
+							LOGGER.warn("{} Cannot parse response: {}", this.name(), completion.response());
+							continue;
+						}
+						
+						parseable = true;
+						break;
+					}
+					catch (final Exception e) {
+						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("☢");
+					return;
+				}
+				
+				if (this.postProcessor != null) {
+					this.postProcessor.accept(split.context());
+				}
+				
+				this.chooseNextAgent(conversation, completion, split.context());
+				
+				retContexts.add(split.context());
+			}, agentExecutorService));
+		}
+		
+		CompletableFuture.allOf(tasks.toArray(new CompletableFuture[0])).join();
+		
+		if (this.merger != null) {
+			final var mergedContexts = this.merger.apply(retContexts);
+			mergedContexts.forEach(c -> c.getGraph().add(retContexts.size() + "↣" + mergedContexts.size()));
+			return mergedContexts;
+		}
+		return retContexts;
+	}
+	
+	public static List<AgentContext> mergeContexts(final List<String> keys, final List<AgentContext> contexts) {
+		final var first = contexts.getFirst();
+		for (final var key : keys) {
+			first.getContext().compute(key, (k, v) -> {
+				final var al = new ArrayList<>();
+				if (v != null) {
+					al.addAll(v);
+				}
+				for (int i = 1; i < contexts.size(); i++) {
+					final var ctx = contexts.get(i);
+					final var val = ctx.getContext().get(key);
+					if (val != null) {
+						al.addAll(val);
+					}
+				}
+				return al;
+			});
+		}
+		return List.of(first);
+	}
+	
+	private List<Split> applyTemplate(final AgentContext context) {
+		String text = this.textTemplate;
+		for (final var entry : context.getContext().entrySet()) {
+			final var k = "{{" + entry.getKey() + "}}";
+			if (text.contains(k)) {
+				text = text.replace(k, entry.getValue().stream().map(Object::toString).collect(Collectors.joining("\n\n")));
+			}
+		}
+		
+		for (final var entry : context.getContext().entrySet()) {
+			final var k = "[[" + entry.getKey() + "]]";
+			if (text.contains(k)) {
+				final var result = new ArrayList<Split>();
+				
+				int i = 0;
+				for (final var value : entry.getValue()) {
+					final var nextText = text.replace(k, value.toString());
+					
+					final var nextContext = new HashMap<>(context.getContext());
+					nextContext.put(entry.getKey(), List.of(value));
+					
+					final var nextAgentContext = new AgentContext(context, nextContext);
+					if (entry.getValue().size() > 1) {
+						nextAgentContext.getGraph().add("⇉ " + ++i + "/" + entry.getValue().size());
+					}
+					
+					result.add(new Split(nextText, nextAgentContext));
+				}
+				
+				return result; // TODO: Multiple splits not supported yet
+			}
+		}
+		
+		return List.of(new Split(text, context));
+	}
+	
+	private Conversation setupConversation(final Split split) {
+		var conversation = split.context().getConversation();
+		if (conversation == null || conversation.getConversation().isEmpty()) {
+			conversation = Conversation.create(this.systemPrompt(), split.text());
+			split.context().setConversation(conversation);
+		}
+		else if (this.agentType() == AgentType.START_CONVERSATION) {
+			conversation = Conversation.create(this.systemPrompt() + "\n" + split.text());
+			split.context().setConversation(conversation);
+		}
+		else {
+			if (conversation.getConversation().get(0).type() == TurnType.SYSTEM) {
+				final var mergedSysAndUser = conversation.getConversation().get(0).text() + "\n" + conversation.getConversation().get(1).text();
+				
+				final var newConversation = Conversation.create(mergedSysAndUser);
+				for (int n = 2; n < conversation.getConversation().size(); n++) {
+					newConversation.addTurn(conversation.getConversation().get(n));
+				}
+				
+				split.context().setConversation(newConversation);
+				conversation = newConversation;
+			}
+			if (conversation.getConversation().get(conversation.getConversation().size() - 1).type() == TurnType.ASSISTANT) {
+				conversation.addTurn(new Conversation.Turn(TurnType.USER, String.join("\n", this.systemPrompt, split.text())));
+			}
+		}
+		return conversation;
+	}
+	
+	private void chooseNextAgent(final Conversation conversation, final Completion completion, final AgentContext context) {
+		context.setNextAgent(this.chooseNext.apply(context));
+		
+		if (context.getNextAgent() != null) {
+			if (context.getNextAgent().gradingPassed() != null) {
+				if (context.getNextAgent().gradingPassed()) {
+					context.getGraph().add("✓");
+				}
+				else {
+					context.getGraph().add("✗");
+				}
+			}
+			
+			if (context.getNextAgent().name() == null) {
+				context.setNextAgent(null);
+				return;
+			}
+			
+			String answer = completion.response();
+			if (context.getNextAgent().transformResponse() != null) {
+				answer = context.getNextAgent().transformResponse().apply(answer);
+			}
+			
+			final Agent nextAgent = context.getAgents().get(context.getNextAgent().name());
+			switch (nextAgent) {
+				case null -> {
+					context.getGraph().add("☢");
+					LOGGER.error("Agent not found: {}", context.getNextAgent().name());
+					context.setNextAgent(null);
+				}
+				
+				case final Agent a when context.getNextAgent().responseAsUserPrompt() == Boolean.TRUE -> {
+					context.getGraph().add("↺");
+					conversation.replaceTurn(new Conversation.Turn(TurnType.USER, answer));
+				}
+				
+				case final Agent a when a.agentType() == AgentType.COMPLETION -> {
+					context.getGraph().add("→");
+					context.setConversation(null);
+				}
+				
+				case final Agent a when a.agentType() == AgentType.START_CONVERSATION -> {
+					context.getGraph().add("🗨");
+					context.setConversation(null);
+				}
+				
+				case final Agent a when a.agentType() == AgentType.CONVERSATION -> {
+					context.getGraph().add("↴");
+					conversation.addTurn(new Conversation.Turn(TurnType.ASSISTANT, answer));
+				}
+				
+				default -> {
+					context.getGraph().add("☢");
+					LOGGER.error("Unknown agent type: {}", nextAgent.agentType());
+					context.setNextAgent(null);
+				}
+			}
+		}
+	}
+}

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

@@ -0,0 +1,78 @@
+package de.extio.lmlib.agent;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+
+import de.extio.lmlib.client.Conversation;
+
+public final class AgentContext {
+	
+	private final Map<String, Agent> agents;
+	
+	private final Map<String, List<Object>> context;
+	
+	private final AgentRequestStatistic requestStatistic;
+	
+	private final List<String> graph = Collections.synchronizedList(new ArrayList<>());
+	
+	private volatile Conversation conversation;
+	
+	private volatile AgentNext nextAgent;
+	
+	public AgentContext(final Map<String, Agent> agents) {
+		this.agents = agents;
+		this.context = new ConcurrentHashMap<>();
+		this.requestStatistic = new AgentRequestStatistic();
+	}
+	
+	public AgentContext(final AgentContext other) {
+		this(other, new ConcurrentHashMap<>(other.context));
+	}
+	
+	public AgentContext(final AgentContext other, final Map<String, List<Object>> context) {
+		this.agents = other.agents;
+		this.context = context;
+		this.conversation = Conversation.copy(other.conversation);
+		this.nextAgent = other.nextAgent;
+		this.requestStatistic = other.requestStatistic;
+		synchronized (other.graph) {
+			this.graph.addAll(other.graph);
+		}
+	}
+	
+	public Map<String, List<Object>> getContext() {
+		return this.context;
+	}
+	
+	public Map<String, Agent> getAgents() {
+		return this.agents;
+	}
+	
+	public AgentRequestStatistic getRequestStatistic() {
+		return this.requestStatistic;
+	}
+	
+	public Conversation getConversation() {
+		return this.conversation;
+	}
+	
+	public void setConversation(final Conversation conversation) {
+		this.conversation = conversation;
+	}
+	
+	public AgentNext getNextAgent() {
+		return this.nextAgent;
+	}
+	
+	public void setNextAgent(final AgentNext next) {
+		this.nextAgent = next;
+	}
+	
+	public List<String> getGraph() {
+		return this.graph;
+	}
+	
+}

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

@@ -0,0 +1,78 @@
+package de.extio.lmlib.agent;
+
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+import org.springframework.beans.factory.DisposableBean;
+import org.springframework.beans.factory.InitializingBean;
+import org.springframework.beans.factory.annotation.Autowired;
+import org.springframework.beans.factory.annotation.Value;
+import org.springframework.stereotype.Service;
+
+import de.extio.lmlib.client.ClientService;
+
+@Service
+public class AgentExecutorService implements InitializingBean, DisposableBean {
+	
+	private static final Logger LOGGER = LoggerFactory.getLogger(AgentExecutorService.class);
+	
+	@Value("${agent.threads:2}")
+	private int threads;
+	
+	@Autowired
+	private ClientService clientService;
+	
+	private ExecutorService branchExecutorService;
+	
+	private ExecutorService agentExecutorService;
+	
+	@Override
+	public void afterPropertiesSet() throws Exception {
+		this.branchExecutorService = Executors.newVirtualThreadPerTaskExecutor();
+		this.agentExecutorService = Executors.newFixedThreadPool(this.threads);
+	}
+	
+	@Override
+	public void destroy() throws Exception {
+		this.branchExecutorService.close();
+		this.agentExecutorService.close();
+	}
+	
+	public List<AgentContext> walkGraph(final Agent agent, final AgentContext context) {
+		LOGGER.info("Agent: {}", agent.name());
+		
+		final var client = this.clientService.getClient(agent.modelCategory());
+		final var responses = Collections.synchronizedList(new ArrayList<AgentContext>());
+		
+		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) {
+				final var branchAgent = context.getAgents().get(branchContext.getNextAgent().name());
+				if (branchAgent == null) {
+					throw new IllegalArgumentException("Agent not found: " + branchContext.getNextAgent().name());
+				}
+				
+				branchTasks.add(CompletableFuture.runAsync(() -> {
+					responses.addAll(this.walkGraph(branchAgent, branchContext));
+				}, this.branchExecutorService));
+			}
+			else {
+				responses.add(branchContext);
+			}
+		}
+		
+		if (!branchTasks.isEmpty()) {
+			CompletableFuture.allOf(branchTasks.toArray(new CompletableFuture[0])).join();
+		}
+		
+		return responses;
+	}
+	
+}

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

@@ -0,0 +1,7 @@
+package de.extio.lmlib.agent;
+
+import java.util.function.Function;
+
+public final record AgentNext(String name, Boolean gradingPassed, Boolean responseAsUserPrompt, Function<String, String> transformResponse) {
+	
+}

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

@@ -0,0 +1,83 @@
+package de.extio.lmlib.agent;
+
+import java.text.DecimalFormat;
+import java.time.Duration;
+import java.time.Instant;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+
+import de.extio.lmlib.client.Completion;
+
+public final class AgentRequestStatistic {
+	
+	private final AtomicInteger requests = new AtomicInteger();
+	
+	private final AtomicLong inTokens = new AtomicLong();
+	
+	private final AtomicLong outTokens = new AtomicLong();
+	
+	private final Instant start = Instant.now();
+	
+	private volatile Duration requestDuration = Duration.ofMillis(0L);
+	
+	public AtomicInteger getRequests() {
+		return this.requests;
+	}
+	
+	public AtomicLong getInTokens() {
+		return this.inTokens;
+	}
+	
+	public AtomicLong getOutTokens() {
+		return this.outTokens;
+	}
+	
+	public Duration getRequestDuration() {
+		return this.requestDuration;
+	}
+	
+	public Duration getAbsoluteDuration() {
+		return Duration.between(this.start, Instant.now());
+	}
+	
+	public String getTps() {
+		return new DecimalFormat("#.##").format((double) (this.inTokens.get() + this.outTokens.get()) / (double)this.getAbsoluteDuration().toMillis() * 1000.0);
+	}
+	
+	public String getOutTps() {
+		return new DecimalFormat("#.##").format((double) this.outTokens.get() / (double)this.getAbsoluteDuration().toMillis() * 1000.0);
+	}
+	
+	public void add(final Completion completion) {
+		this.getRequests().addAndGet(completion.requests());
+		this.getInTokens().addAndGet(completion.inTokens());
+		this.getOutTokens().addAndGet(completion.outTokens());
+		this.addRequestDuration(completion.duration());
+	}
+	
+	public synchronized void addRequestDuration(final Duration duration) {
+		this.requestDuration = this.requestDuration.plus(duration);
+	}
+	
+	@Override
+	public String toString() {
+		final StringBuilder builder = new StringBuilder();
+		builder.append("AgentRequestStatistic [requests=");
+		builder.append(this.requests);
+		builder.append(", inTokens=");
+		builder.append(this.inTokens);
+		builder.append(", outTokens=");
+		builder.append(this.outTokens);
+		builder.append(", requestDuration=");
+		builder.append(this.requestDuration);
+		builder.append(", getAbsoluteDuration()=");
+		builder.append(this.getAbsoluteDuration());
+		builder.append(", getTps()=");
+		builder.append(this.getTps());
+		builder.append(", getOutTps()=");
+		builder.append(this.getOutTps());
+		builder.append("]");
+		return builder.toString();
+	}
+	
+}

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

@@ -0,0 +1,9 @@
+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/AgentType.java

@@ -0,0 +1,5 @@
+package de.extio.lmlib.agent;
+
+public enum AgentType {
+	COMPLETION, START_CONVERSATION, CONVERSATION;
+}

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

@@ -0,0 +1,57 @@
+package de.extio.lmlib.agent;
+
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.List;
+import java.util.Map;
+
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.fasterxml.jackson.databind.JsonNode;
+import com.fasterxml.jackson.databind.ObjectMapper;
+
+import de.extio.lmlib.client.Completion;
+
+public class JsonAgentResponseHandler implements AgentResponseHandler {
+	
+	private static final Logger LOGGER = LoggerFactory.getLogger(JsonAgentResponseHandler.class);
+	
+	private final ObjectMapper objectMapper = new ObjectMapper();
+	
+	@Override
+	public boolean handle(final Split split, final Completion completion) {
+		final int firstCurly = completion.response().indexOf("{");
+		final int lastCurly = completion.response().lastIndexOf("}");
+		final var response = completion.response().substring(firstCurly, lastCurly + 1);
+		JsonNode jsonNode;
+		try {
+			jsonNode = this.objectMapper.readTree(response);
+		}
+		catch (final Exception e) {
+			LOGGER.warn("Cannot parse json response: {}", completion.response());
+			return false;
+		}
+		
+		boolean hasResponse = false;
+		final Iterator<Map.Entry<String, JsonNode>> fields = jsonNode.fields();
+		while (fields.hasNext()) {
+			final var entry = fields.next();
+			if (entry.getValue().isArray()) {
+				final var items = new HashSet<>();
+				entry.getValue().elements().forEachRemaining(item -> items.add(item.asText()));
+				split.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()));
+				hasResponse = true;
+				LOGGER.debug("Response: {} {}", entry.getKey(), entry.getValue().asText());
+			}
+		}
+		
+		return hasResponse;
+	}
+	
+}

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

@@ -0,0 +1,5 @@
+package de.extio.lmlib.agent;
+
+record Split(String text, AgentContext context) {
+	
+}

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

@@ -0,0 +1,27 @@
+package de.extio.lmlib.agent;
+
+import java.util.List;
+
+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;
+	
+	public TextAgentResponseHandler(final String key) {
+		this.key = key;
+	}
+	
+	@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());
+		return true;
+	}
+	
+}