refactor: Update modelCategory method to accept AgentContext and adjust related calls
This commit is contained in:
@@ -36,6 +36,14 @@ public record Agent(String name,
|
|||||||
return BaseAgent.super.agentType(context);
|
return BaseAgent.super.agentType(context);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
@Override
|
||||||
|
public ModelCategory modelCategory(final AgentContext context) {
|
||||||
|
if (modelCategory != null) {
|
||||||
|
return modelCategory;
|
||||||
|
}
|
||||||
|
return BaseAgent.super.modelCategory(context);
|
||||||
|
}
|
||||||
|
|
||||||
@Override
|
@Override
|
||||||
public void preProcess(final AgentContext context) {
|
public void preProcess(final AgentContext context) {
|
||||||
if (preProcessor != null) {
|
if (preProcessor != null) {
|
||||||
|
|||||||
@@ -68,7 +68,7 @@ public class AgentExecutorService implements InitializingBean, DisposableBean {
|
|||||||
|
|
||||||
Client client = null;
|
Client client = null;
|
||||||
if (agent.agentType(context) != AgentType.PROCESSING_ONLY) {
|
if (agent.agentType(context) != AgentType.PROCESSING_ONLY) {
|
||||||
client = this.clientService.getClient(agent.modelCategory());
|
client = this.clientService.getClient(agent.modelCategory(context));
|
||||||
}
|
}
|
||||||
context.setAgentContextUpdateConsumer(agentContextUpdateConsumer);
|
context.setAgentContextUpdateConsumer(agentContextUpdateConsumer);
|
||||||
|
|
||||||
|
|||||||
@@ -34,7 +34,11 @@ public interface BaseAgent {
|
|||||||
throw new UnsupportedOperationException("Name must be implemented in the agent");
|
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;
|
return ModelCategory.MEDIUM;
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -50,10 +54,6 @@ public interface BaseAgent {
|
|||||||
return new TextAgentResponseHandler("response");
|
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) {
|
default void preProcess(final AgentContext context) {
|
||||||
|
|
||||||
}
|
}
|
||||||
@@ -88,7 +88,7 @@ public interface BaseAgent {
|
|||||||
split.context().getGraph().add("○");
|
split.context().getGraph().add("○");
|
||||||
}
|
}
|
||||||
else {
|
else {
|
||||||
split.context().getGraph().add(this.modelCategory().getShortName());
|
split.context().getGraph().add(this.modelCategory(split.context()).getShortName());
|
||||||
}
|
}
|
||||||
split.context().getGraph().add(this.name());
|
split.context().getGraph().add(this.name());
|
||||||
|
|
||||||
@@ -103,7 +103,7 @@ public interface BaseAgent {
|
|||||||
if (this.responseHandler() instanceof final StreamedAgentResponseHandler streamedResponseHandler) {
|
if (this.responseHandler() instanceof final StreamedAgentResponseHandler streamedResponseHandler) {
|
||||||
streamedResponseHandler.beforeStream(split.context());
|
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 (this.responseHandler() instanceof final StreamedAgentResponseHandler streamedResponseHandler) {
|
||||||
if (streamedResponseHandler.handleChunk(chunk, split.context()) && split.context().getAgentContextUpdateConsumer() != null) {
|
if (streamedResponseHandler.handleChunk(chunk, split.context()) && split.context().getAgentContextUpdateConsumer() != null) {
|
||||||
split.context().getAgentContextUpdateConsumer().accept(split.context());
|
split.context().getAgentContextUpdateConsumer().accept(split.context());
|
||||||
@@ -112,7 +112,7 @@ public interface BaseAgent {
|
|||||||
});
|
});
|
||||||
}
|
}
|
||||||
else {
|
else {
|
||||||
completion = client.conversation(this.modelCategory(), conversation);
|
completion = client.conversation(this.modelCategory(split.context()), conversation);
|
||||||
}
|
}
|
||||||
split.context().setLastCompletion(completion);
|
split.context().setLastCompletion(completion);
|
||||||
split.context().getRequestStatistic().add(completion);
|
split.context().getRequestStatistic().add(completion);
|
||||||
|
|||||||
Reference in New Issue
Block a user