From a176e6b6c4fc7a93b5a6eb7e62c6b5b91a31bea3 Mon Sep 17 00:00:00 2001 From: extio-de Date: Sun, 13 Jul 2025 10:42:02 +0200 Subject: [PATCH] refactor: Update modelCategory method to accept AgentContext and adjust related calls --- src/main/java/de/extio/lmlib/agent/Agent.java | 8 ++++++++ .../lmlib/agent/AgentExecutorService.java | 2 +- .../java/de/extio/lmlib/agent/BaseAgent.java | 18 +++++++++--------- 3 files changed, 18 insertions(+), 10 deletions(-) diff --git a/src/main/java/de/extio/lmlib/agent/Agent.java b/src/main/java/de/extio/lmlib/agent/Agent.java index 7724214..e1c1d44 100755 --- a/src/main/java/de/extio/lmlib/agent/Agent.java +++ b/src/main/java/de/extio/lmlib/agent/Agent.java @@ -36,6 +36,14 @@ public record Agent(String name, return BaseAgent.super.agentType(context); } + @Override + public ModelCategory modelCategory(final AgentContext context) { + if (modelCategory != null) { + return modelCategory; + } + return BaseAgent.super.modelCategory(context); + } + @Override public void preProcess(final AgentContext context) { if (preProcessor != null) { diff --git a/src/main/java/de/extio/lmlib/agent/AgentExecutorService.java b/src/main/java/de/extio/lmlib/agent/AgentExecutorService.java index 175a284..df61819 100755 --- a/src/main/java/de/extio/lmlib/agent/AgentExecutorService.java +++ b/src/main/java/de/extio/lmlib/agent/AgentExecutorService.java @@ -68,7 +68,7 @@ public class AgentExecutorService implements InitializingBean, DisposableBean { Client client = null; if (agent.agentType(context) != AgentType.PROCESSING_ONLY) { - client = this.clientService.getClient(agent.modelCategory()); + client = this.clientService.getClient(agent.modelCategory(context)); } context.setAgentContextUpdateConsumer(agentContextUpdateConsumer); diff --git a/src/main/java/de/extio/lmlib/agent/BaseAgent.java b/src/main/java/de/extio/lmlib/agent/BaseAgent.java index 50b114b..8330487 100644 --- a/src/main/java/de/extio/lmlib/agent/BaseAgent.java +++ b/src/main/java/de/extio/lmlib/agent/BaseAgent.java @@ -34,7 +34,11 @@ public interface BaseAgent { throw new UnsupportedOperationException("Name must be implemented in the agent"); } - default ModelCategory modelCategory() { + default AgentType agentType(final AgentContext context) { + throw new UnsupportedOperationException("Agent type must be implemented in the agent"); + } + + default ModelCategory modelCategory(final AgentContext context) { return ModelCategory.MEDIUM; } @@ -49,10 +53,6 @@ public interface BaseAgent { default AgentResponseHandler responseHandler() { return new TextAgentResponseHandler("response"); } - - default AgentType agentType(final AgentContext context) { - throw new UnsupportedOperationException("Agent type must be implemented in the agent"); - } default void preProcess(final AgentContext context) { @@ -88,7 +88,7 @@ public interface BaseAgent { split.context().getGraph().add("○"); } else { - split.context().getGraph().add(this.modelCategory().getShortName()); + split.context().getGraph().add(this.modelCategory(split.context()).getShortName()); } split.context().getGraph().add(this.name()); @@ -102,8 +102,8 @@ public interface BaseAgent { if (split.context().isStreaming()) { if (this.responseHandler() instanceof final StreamedAgentResponseHandler streamedResponseHandler) { streamedResponseHandler.beforeStream(split.context()); - } - completion = client.streamConversation(this.modelCategory(), conversation, chunk -> { + } + completion = client.streamConversation(this.modelCategory(split.context()), 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()); @@ -112,7 +112,7 @@ public interface BaseAgent { }); } else { - completion = client.conversation(this.modelCategory(), conversation); + completion = client.conversation(this.modelCategory(split.context()), conversation); } split.context().setLastCompletion(completion); split.context().getRequestStatistic().add(completion);