Skip to content

feat(core): support per-session runtime model switching #2097

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Buktal wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat(core): support per-session runtime model switching #2097

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
106 changes: 94 additions & 12 deletions agentscope-core/src/main/java/io/agentscope/core/ReActAgent.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2118,7 +2118,7 @@ Flux<AgentEvent> reasoningStream(
rc,
MiddlewareBase::onModelCall,
modelCallCore)
.apply(new ModelCallInput(messages, tools, options, modelForCall()))
.apply(new ModelCallInput(messages, tools, options, modelForCall(rc)))
.doOnNext(this::publishEvent);
}

Expand Down Expand Up @@ -3052,7 +3052,7 @@ protected Mono<Msg> summarizing() {
List<Msg> messageList = prepareSummaryMessages();
GenerateOptions generateOptions = buildGenerateOptions();
ReasoningContext context = new ReasoningContext(getName());
Model summaryModel = modelForCall();
Model summaryModel = modelForCall(rc);
publishEvent(new ExceedMaxItersEvent("", maxIters, maxIters));

return hookDispatcher
Expand Down Expand Up @@ -3482,29 +3482,63 @@ protected GenerateOptions buildGenerateOptions() {
return baseOptions != null ? baseOptions : GenerateOptions.builder().build();
}

private Model modelForCall() {
Model fallbackModel = modelConfig.fallbackModel();
if (fallbackModel == null) {
return model;
/**
* Resolves the model for a single call, wrapping it with the configured fallback when one is
* present.
*/
private Model modelForCall(RuntimeContext ctx) {
Model primary = getModel(ctx);
Model fallback = modelConfig.fallbackModel();
return fallback == null ? primary : wrapWithFallback(primary, fallback);
}

/**
* Returns the model effective for the session carried by {@code ctx}: the session's registered
* id resolved through {@link ModelRegistry} when set, otherwise the agent's default model. An
* unresolvable session id is logged and falls back to the default.
*
* @param ctx the runtime context identifying the session (may be {@code null})
* @return the resolved session model, or the default model when no id is set or it fails to
* resolve
*/
public Model getModel(RuntimeContext ctx) {
AgentState s = ctx != null ? ctx.getAgentState() : null;
String mid = s != null ? s.getModelId() : null;
if (mid != null && !mid.isBlank()) {
try {
return ModelRegistry.resolve(mid);
} catch (Exception e) {
log.warn(
"Session modelId [{}] unresolvable, fallback to default model: {}",
mid,
e.getMessage());
}
}
return model;
}

AtomicReference<Model> activeModel = new AtomicReference<>(model);
/**
* Wraps {@code primary} so that {@code fallback} takes over if the first signal from
* {@code primary} is an error.
*/
private static Model wrapWithFallback(Model primary, Model fallback) {
AtomicReference<Model> activeModel = new AtomicReference<>(primary);
return new Model() {
@Override
public Flux<ChatResponse> stream(
List<Msg> messages, List<ToolSchema> tools, GenerateOptions options) {
Flux<ChatResponse> primaryFlux = model.stream(messages, tools, options);
Flux<ChatResponse> primaryFlux = primary.stream(messages, tools, options);
return primaryFlux.switchOnFirst(
(signal, flux) -> {
if (signal.isOnError()) {
Throwable error = signal.getThrowable();
activeModel.set(fallbackModel);
activeModel.set(fallback);
log.warn(
"Primary model {} failed, switching to fallback {}",
model.getModelName(),
fallbackModel.getModelName(),
primary.getModelName(),
fallback.getModelName(),
error);
return fallbackModel.stream(messages, tools, options);
return fallback.stream(messages, tools, options);
}
return flux;
});
Expand Down Expand Up @@ -3693,6 +3727,54 @@ public void setPermissionMode(RuntimeContext ctx, PermissionMode mode) {
setPermissionMode(uid, sid, mode);
}

/**
* Switches the model used by the given {@code (userId, sessionId)} session at runtime by storing
* its id in the persisted {@link AgentState}. The id is resolved through {@link ModelRegistry}
* on each {@code call}, so a blank or {@code null} id restores the agent's default model. The
* change is persisted so the next {@code call} on that session—and any cluster node that
* reloads the state—sees it.
*
* <p>Unlike {@link #setPermissionMode(String, String, PermissionMode)}, no per-slot cache is
* rebuilt: models are stateless, so resolution happens fresh on every call.
*
* @param userId user identity for the slot (may be {@code null})
* @param sessionId session identity (falls back to the default session id when {@code null})
* @param modelId the registered model id to use, or {@code null}/blank to restore the default
*/
public void setModel(String userId, String sessionId, String modelId) {
String sid = (sessionId == null || sessionId.isBlank()) ? defaultSessionId : sessionId;
AgentState s = getAgentState(userId, sid);
s.setModelId(modelId);
saveAgentState(userId, sid);
}

/**
* Switches the model for the session identified by the given {@link RuntimeContext}. See
* {@link #setModel(String, String, String)}.
*
* @param ctx the runtime context identifying the session
* @param modelId the registered model id to use, or {@code null}/blank to restore the default
*/
public void setModel(RuntimeContext ctx, String modelId) {
String uid = ctx != null ? ctx.getUserId() : null;
String sid = ctx != null ? ctx.getSessionId() : null;
setModel(uid, sid, modelId);
}

/**
* Clears the session's model id so subsequent calls revert to the agent's default model. The
* change is persisted.
*
* @param userId user identity for the slot (may be {@code null})
* @param sessionId session identity (falls back to the default session id when {@code null})
*/
public void clearModel(String userId, String sessionId) {
String sid = (sessionId == null || sessionId.isBlank()) ? defaultSessionId : sessionId;
AgentState s = getAgentState(userId, sid);
s.setModelId(null);
saveAgentState(userId, sid);
}

/**
* Returns the current {@link PermissionMode} for the given {@code (userId, sessionId)} session.
*
Expand Down
32 changes: 28 additions & 4 deletions agentscope-core/src/main/java/io/agentscope/core/state/AgentState.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,8 @@
"permission_context",
"tool_context",
"tasks_context",
"plan_mode_context"
"plan_mode_context",
"model_id"
})
public final class AgentState implements State {

Expand All @@ -69,6 +70,7 @@ public final class AgentState implements State {
private final ToolContextState toolContext;
private final TaskContextState tasksContext;
private final PlanModeContextState planModeContext;
private String modelId;

/**
* Per-session interrupt signal. Runtime-only (never serialized): a stateless agent engine
Expand Down Expand Up @@ -100,6 +102,7 @@ private AgentState(Builder builder) {
builder.planModeContext == null
? new PlanModeContextState()
: builder.planModeContext;
this.modelId = builder.modelId;
}

@JsonCreator
Expand All @@ -114,7 +117,8 @@ static AgentState fromJson(
@JsonProperty("permission_context") PermissionContextState permissionContext,
@JsonProperty("tool_context") ToolContextState toolContext,
@JsonProperty("tasks_context") TaskContextState tasksContext,
@JsonProperty("plan_mode_context") PlanModeContextState planModeContext) {
@JsonProperty("plan_mode_context") PlanModeContextState planModeContext,
@JsonProperty("model_id") String modelId) {
Builder b = builder();
if (sessionId != null) {
b.sessionId(sessionId);
Expand Down Expand Up @@ -149,6 +153,9 @@ static AgentState fromJson(
if (planModeContext != null) {
b.planModeContext(planModeContext);
}
if (modelId != null) {
b.modelId(modelId);
}
return b.build();
}

Expand Down Expand Up @@ -247,6 +254,15 @@ public PlanModeContextState getPlanModeContext() {
return planModeContext;
}

@JsonProperty("model_id")
public String getModelId() {
return modelId;
}

public void setModelId(String modelId) {
this.modelId = (modelId == null || modelId.isBlank()) ? null : modelId;
}

/**
* The per-session interrupt signal for this state, created on first access. Runtime-only and not
* serialized.
Expand Down Expand Up @@ -305,7 +321,8 @@ public boolean equals(Object o) {
&& Objects.equals(permissionContext, other.permissionContext)
&& Objects.equals(toolContext, other.toolContext)
&& Objects.equals(tasksContext, other.tasksContext)
&& Objects.equals(planModeContext, other.planModeContext);
&& Objects.equals(planModeContext, other.planModeContext)
&& Objects.equals(modelId, other.modelId);
}

@Override
Expand All @@ -321,7 +338,8 @@ public int hashCode() {
permissionContext,
toolContext,
tasksContext,
planModeContext);
planModeContext,
modelId);
}

@Override
Expand Down Expand Up @@ -351,6 +369,7 @@ public static final class Builder {
private ToolContextState toolContext;
private TaskContextState tasksContext;
private PlanModeContextState planModeContext;
private String modelId;

private Builder() {}

Expand Down Expand Up @@ -416,6 +435,11 @@ public Builder planModeContext(PlanModeContextState planModeContext) {
return this;
}

public Builder modelId(String modelId) {
this.modelId = modelId;
return this;
}

public AgentState build() {
return new AgentState(this);
}
Expand Down
Loading
Loading