Skip to content

Commit bb47787

Browse files
mingshlZhangxunmt
mingshl
and
Zhangxunmt
authored
[Feature branch] Introduce hook and context management to OpenSearch Agents (#4397)
* add hooks in ml-commons (#4326) Signed-off-by: Xun Zhang <[email protected]> * initiate context management api with hook implementation (#4345) * initiate context management api with hook implementation Signed-off-by: Mingshi Liu <[email protected]> * apply spotless Signed-off-by: Mingshi Liu <[email protected]> --------- Signed-off-by: Mingshi Liu <[email protected]> * Add Context Manager to PER (#4379) * add pre_llm hook to per agent Signed-off-by: Mingshi Liu <[email protected]> change context management passing from query parameters to payload Signed-off-by: Mingshi Liu <[email protected]> pass hook registery into PER Signed-off-by: Mingshi Liu <[email protected]> apply spotless Signed-off-by: Mingshi Liu <[email protected]> initiate context management api with hook implementation Signed-off-by: Mingshi Liu <[email protected]> * add comment Signed-off-by: Mingshi Liu <[email protected]> * format Signed-off-by: Mingshi Liu <[email protected]> * add validation Signed-off-by: Mingshi Liu <[email protected]> --------- Signed-off-by: Mingshi Liu <[email protected]> * add inner create context management to agent register api Signed-off-by: Mingshi Liu <[email protected]> * add code coverage Signed-off-by: Mingshi Liu <[email protected]> * allow context management hook register in during agent execute Signed-off-by: Mingshi Liu <[email protected]> * add code coverage Signed-off-by: Mingshi Liu <[email protected]> * add more code coverage Signed-off-by: Mingshi Liu <[email protected]> * add validation check Signed-off-by: Mingshi Liu <[email protected]> * adapt to inplace update for context Signed-off-by: Mingshi Liu <[email protected]> * fix test Signed-off-by: Mingshi Liu <[email protected]> --------- Signed-off-by: Xun Zhang <[email protected]> Signed-off-by: Mingshi Liu <[email protected]> Co-authored-by: Xun Zhang <[email protected]>
1 parent d66e924 commit bb47787

File tree

89 files changed

+9799
-1555
lines changed

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

89 files changed

+9799
-1555
lines changed

‎common/src/main/java/org/opensearch/ml/common/CommonValue.java‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,7 @@ public class CommonValue {
5454
public static final String TASK_POLLING_JOB_INDEX = ".ml_commons_task_polling_job";
5555
public static final String MCP_SESSION_MANAGEMENT_INDEX = ".plugins-ml-mcp-session-management";
5656
public static final String MCP_TOOLS_INDEX = ".plugins-ml-mcp-tools";
57+
public static final String ML_CONTEXT_MANAGEMENT_TEMPLATES_INDEX = ".plugins-ml-context-management-templates";
5758
// index created in 3.1 to track all ml jobs created via job scheduler
5859
public static final String ML_JOBS_INDEX = ".plugins-ml-jobs";
5960
public static final Set<String> stopWordsIndices = ImmutableSet.of(".plugins-ml-stop-words");
@@ -76,6 +77,7 @@ public class CommonValue {
7677
public static final String ML_LONG_MEMORY_HISTORY_INDEX_MAPPING_PATH = "index-mappings/ml_memory_long_term_history.json";
7778
public static final String ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_session_management.json";
7879
public static final String ML_MCP_TOOLS_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_tools.json";
80+
public static final String ML_CONTEXT_MANAGEMENT_TEMPLATES_INDEX_MAPPING_PATH = "index-mappings/ml_context_management_templates.json";
7981
public static final String ML_JOBS_INDEX_MAPPING_PATH = "index-mappings/ml_jobs.json";
8082
public static final String ML_INDEX_INSIGHT_CONFIG_INDEX_MAPPING_PATH = "index-mappings/ml_index_insight_config.json";
8183
public static final String ML_INDEX_INSIGHT_STORAGE_INDEX_MAPPING_PATH = "index-mappings/ml_index_insight_storage.json";

‎common/src/main/java/org/opensearch/ml/common/agent/MLAgent.java‎

Lines changed: 104 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
import org.opensearch.ml.common.CommonValue;
3131
import org.opensearch.ml.common.MLAgentType;
3232
import org.opensearch.ml.common.MLModel;
33+
import org.opensearch.ml.common.contextmanager.ContextManagementTemplate;
3334
import org.opensearch.telemetry.metrics.tags.Tags;
3435

3536
import lombok.Builder;
@@ -52,13 +53,16 @@ public class MLAgent implements ToXContentObject, Writeable {
5253
public static final String LAST_UPDATED_TIME_FIELD = "last_updated_time";
5354
public static final String APP_TYPE_FIELD = "app_type";
5455
public static final String IS_HIDDEN_FIELD = "is_hidden";
56+
public static final String CONTEXT_MANAGEMENT_NAME_FIELD = "context_management_name";
57+
public static final String CONTEXT_MANAGEMENT_FIELD = "context_management";
5558
private static final String LLM_INTERFACE_FIELD = "_llm_interface";
5659
private static final String TAG_VALUE_UNKNOWN = "unknown";
5760
private static final String TAG_MEMORY_TYPE = "memory_type";
5861

5962
public static final int AGENT_NAME_MAX_LENGTH = 128;
6063

6164
private static final Version MINIMAL_SUPPORTED_VERSION_FOR_HIDDEN_AGENT = CommonValue.VERSION_2_13_0;
65+
private static final Version MINIMAL_SUPPORTED_VERSION_FOR_CONTEXT_MANAGEMENT = CommonValue.VERSION_3_3_0;
6266

6367
private String name;
6468
private String type;
@@ -73,6 +77,8 @@ public class MLAgent implements ToXContentObject, Writeable {
7377
private Instant lastUpdateTime;
7478
private String appType;
7579
private Boolean isHidden;
80+
private String contextManagementName;
81+
private ContextManagementTemplate contextManagement;
7682
private final String tenantId;
7783

7884
@Builder(toBuilder = true)
@@ -89,6 +95,8 @@ public MLAgent(
8995
Instant lastUpdateTime,
9096
String appType,
9197
Boolean isHidden,
98+
String contextManagementName,
99+
ContextManagementTemplate contextManagement,
92100
String tenantId
93101
) {
94102
this.name = name;
@@ -104,6 +112,8 @@ public MLAgent(
104112
this.appType = appType;
105113
// is_hidden field isn't going to be set by user. It will be set by the code.
106114
this.isHidden = isHidden;
115+
this.contextManagementName = contextManagementName;
116+
this.contextManagement = contextManagement;
107117
this.tenantId = tenantId;
108118
validate();
109119
}
@@ -123,7 +133,23 @@ public MLAgent(
123133
Boolean isHidden,
124134
String tenantId
125135
) {
126-
this(name, type, description, llm, null, tools, parameters, memory, createdTime, lastUpdateTime, appType, isHidden, tenantId);
136+
this(
137+
name,
138+
type,
139+
description,
140+
llm,
141+
null,
142+
tools,
143+
parameters,
144+
memory,
145+
createdTime,
146+
lastUpdateTime,
147+
appType,
148+
isHidden,
149+
null,
150+
null,
151+
tenantId
152+
);
127153
}
128154

129155
private void validate() {
@@ -150,6 +176,17 @@ private void validate() {
150176
}
151177
}
152178
}
179+
validateContextManagement();
180+
}
181+
182+
private void validateContextManagement() {
183+
if (contextManagementName != null && contextManagement != null) {
184+
throw new IllegalArgumentException("Cannot specify both context_management_name and context_management");
185+
}
186+
187+
if (contextManagement != null && !contextManagement.isValid()) {
188+
throw new IllegalArgumentException("Invalid context management configuration");
189+
}
153190
}
154191

155192
private void validateMLAgentType(String agentType) {
@@ -196,6 +233,12 @@ public MLAgent(StreamInput input) throws IOException {
196233
if (streamInputVersion.onOrAfter(MINIMAL_SUPPORTED_VERSION_FOR_HIDDEN_AGENT)) {
197234
isHidden = input.readOptionalBoolean();
198235
}
236+
if (streamInputVersion.onOrAfter(MINIMAL_SUPPORTED_VERSION_FOR_CONTEXT_MANAGEMENT)) {
237+
contextManagementName = input.readOptionalString();
238+
if (input.readBoolean()) {
239+
contextManagement = new ContextManagementTemplate(input);
240+
}
241+
}
199242
this.tenantId = streamInputVersion.onOrAfter(VERSION_2_19_0) ? input.readOptionalString() : null;
200243
validate();
201244
}
@@ -245,6 +288,15 @@ public void writeTo(StreamOutput out) throws IOException {
245288
if (streamOutputVersion.onOrAfter(MINIMAL_SUPPORTED_VERSION_FOR_HIDDEN_AGENT)) {
246289
out.writeOptionalBoolean(isHidden);
247290
}
291+
if (streamOutputVersion.onOrAfter(MINIMAL_SUPPORTED_VERSION_FOR_CONTEXT_MANAGEMENT)) {
292+
out.writeOptionalString(contextManagementName);
293+
if (contextManagement != null) {
294+
out.writeBoolean(true);
295+
contextManagement.writeTo(out);
296+
} else {
297+
out.writeBoolean(false);
298+
}
299+
}
248300
if (streamOutputVersion.onOrAfter(VERSION_2_19_0)) {
249301
out.writeOptionalString(tenantId);
250302
}
@@ -290,6 +342,12 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
290342
if (isHidden != null) {
291343
builder.field(MLModel.IS_HIDDEN_FIELD, isHidden);
292344
}
345+
if (contextManagementName != null) {
346+
builder.field(CONTEXT_MANAGEMENT_NAME_FIELD, contextManagementName);
347+
}
348+
if (contextManagement != null) {
349+
builder.field(CONTEXT_MANAGEMENT_FIELD, contextManagement);
350+
}
293351
if (tenantId != null) {
294352
builder.field(TENANT_ID_FIELD, tenantId);
295353
}
@@ -318,6 +376,8 @@ private static MLAgent parseCommonFields(XContentParser parser, boolean parseHid
318376
Instant lastUpdateTime = null;
319377
String appType = null;
320378
boolean isHidden = false;
379+
String contextManagementName = null;
380+
ContextManagementTemplate contextManagement = null;
321381
String tenantId = null;
322382

323383
ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
@@ -367,6 +427,12 @@ private static MLAgent parseCommonFields(XContentParser parser, boolean parseHid
367427
if (parseHidden)
368428
isHidden = parser.booleanValue();
369429
break;
430+
case CONTEXT_MANAGEMENT_NAME_FIELD:
431+
contextManagementName = parser.text();
432+
break;
433+
case CONTEXT_MANAGEMENT_FIELD:
434+
contextManagement = ContextManagementTemplate.parse(parser);
435+
break;
370436
case TENANT_ID_FIELD:
371437
tenantId = parser.textOrNull();
372438
break;
@@ -390,6 +456,8 @@ private static MLAgent parseCommonFields(XContentParser parser, boolean parseHid
390456
.lastUpdateTime(lastUpdateTime)
391457
.appType(appType)
392458
.isHidden(isHidden)
459+
.contextManagementName(contextManagementName)
460+
.contextManagement(contextManagement)
393461
.tenantId(tenantId)
394462
.build();
395463
}
@@ -423,4 +491,39 @@ public Tags getTags() {
423491

424492
return tags;
425493
}
494+
495+
/**
496+
* Check if this agent has context management configuration
497+
* @return true if agent has either context management name or inline configuration
498+
*/
499+
public boolean hasContextManagement() {
500+
return contextManagementName != null || contextManagement != null;
501+
}
502+
503+
/**
504+
* Get the effective context management configuration for this agent.
505+
* This method prioritizes inline configuration over template reference.
506+
* Note: Template resolution requires external service call and should be handled by the caller.
507+
*
508+
* @return the inline context management configuration, or null if using template reference or no configuration
509+
*/
510+
public ContextManagementTemplate getInlineContextManagement() {
511+
return contextManagement;
512+
}
513+
514+
/**
515+
* Check if this agent uses a context management template reference
516+
* @return true if agent references a context management template by name
517+
*/
518+
public boolean hasContextManagementTemplate() {
519+
return contextManagementName != null;
520+
}
521+
522+
/**
523+
* Get the context management template name if this agent references one
524+
* @return the template name, or null if no template reference
525+
*/
526+
public String getContextManagementTemplateName() {
527+
return contextManagementName;
528+
}
426529
}

‎common/src/main/java/org/opensearch/ml/common/contextmanager/ActivationRule.java‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.contextmanager;
7+
8+
/**
9+
* Interface for activation rules that determine when a context manager should execute.
10+
* Activation rules evaluate runtime conditions based on the current context state.
11+
*/
12+
public interface ActivationRule {
13+
14+
/**
15+
* Evaluate whether the activation condition is met.
16+
* @param context the current context state
17+
* @return true if the condition is met and the manager should activate, false otherwise
18+
*/
19+
boolean evaluate(ContextManagerContext context);
20+
21+
/**
22+
* Get a description of this activation rule for logging and debugging.
23+
* @return a human-readable description of the rule
24+
*/
25+
String getDescription();
26+
}

‎common/src/main/java/org/opensearch/ml/common/contextmanager/ActivationRuleFactory.java‎

Lines changed: 146 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.common.contextmanager;
7+
8+
import java.util.ArrayList;
9+
import java.util.List;
10+
import java.util.Map;
11+
12+
import lombok.extern.log4j.Log4j2;
13+
14+
/**
15+
* Factory class for creating activation rules from configuration.
16+
* Supports creating rules from configuration maps and combining multiple rules.
17+
*/
18+
@Log4j2
19+
public class ActivationRuleFactory {
20+
21+
public static final String TOKENS_EXCEED_KEY = "tokens_exceed";
22+
public static final String MESSAGE_COUNT_EXCEED_KEY = "message_count_exceed";
23+
24+
/**
25+
* Create activation rules from a configuration map.
26+
* @param activationConfig the configuration map containing rule definitions
27+
* @return a list of activation rules, or empty list if no valid rules found
28+
*/
29+
public static List<ActivationRule> createRules(Map<String, Object> activationConfig) {
30+
List<ActivationRule> rules = new ArrayList<>();
31+
32+
if (activationConfig == null || activationConfig.isEmpty()) {
33+
return rules;
34+
}
35+
36+
// Create tokens_exceed rule
37+
if (activationConfig.containsKey(TOKENS_EXCEED_KEY)) {
38+
try {
39+
Object tokenValue = activationConfig.get(TOKENS_EXCEED_KEY);
40+
int tokenThreshold = parseIntegerValue(tokenValue, TOKENS_EXCEED_KEY);
41+
if (tokenThreshold > 0) {
42+
rules.add(new TokensExceedRule(tokenThreshold));
43+
log.debug("Created TokensExceedRule with threshold: {}", tokenThreshold);
44+
} else {
45+
log.warn("Invalid token threshold value: {}. Must be positive integer.", tokenValue);
46+
}
47+
} catch (Exception e) {
48+
log.error("Failed to create TokensExceedRule: {}", e.getMessage());
49+
}
50+
}
51+
52+
// Create message_count_exceed rule
53+
if (activationConfig.containsKey(MESSAGE_COUNT_EXCEED_KEY)) {
54+
try {
55+
Object messageValue = activationConfig.get(MESSAGE_COUNT_EXCEED_KEY);
56+
int messageThreshold = parseIntegerValue(messageValue, MESSAGE_COUNT_EXCEED_KEY);
57+
if (messageThreshold > 0) {
58+
rules.add(new MessageCountExceedRule(messageThreshold));
59+
log.debug("Created MessageCountExceedRule with threshold: {}", messageThreshold);
60+
} else {
61+
log.warn("Invalid message count threshold value: {}. Must be positive integer.", messageValue);
62+
}
63+
} catch (Exception e) {
64+
log.error("Failed to create MessageCountExceedRule: {}", e.getMessage());
65+
}
66+
}
67+
68+
return rules;
69+
}
70+
71+
/**
72+
* Create a composite rule that requires ALL rules to be satisfied (AND logic).
73+
* @param rules the list of rules to combine
74+
* @return a composite rule, or null if the list is empty
75+
*/
76+
public static ActivationRule createCompositeRule(List<ActivationRule> rules) {
77+
if (rules == null || rules.isEmpty()) {
78+
return null;
79+
}
80+
81+
if (rules.size() == 1) {
82+
return rules.get(0);
83+
}
84+
85+
return new CompositeActivationRule(rules);
86+
}
87+
88+
/**
89+
* Parse an integer value from configuration, handling various input types.
90+
* @param value the value to parse
91+
* @param fieldName the field name for error reporting
92+
* @return the parsed integer value
93+
* @throws IllegalArgumentException if the value cannot be parsed
94+
*/
95+
private static int parseIntegerValue(Object value, String fieldName) {
96+
if (value instanceof Integer) {
97+
return (Integer) value;
98+
} else if (value instanceof Number) {
99+
return ((Number) value).intValue();
100+
} else if (value instanceof String) {
101+
try {
102+
return Integer.parseInt((String) value);
103+
} catch (NumberFormatException e) {
104+
throw new IllegalArgumentException("Invalid integer value for " + fieldName + ": " + value);
105+
}
106+
} else {
107+
throw new IllegalArgumentException("Unsupported value type for " + fieldName + ": " + value.getClass().getSimpleName());
108+
}
109+
}
110+
111+
/**
112+
* Composite activation rule that implements AND logic for multiple rules.
113+
*/
114+
private static class CompositeActivationRule implements ActivationRule {
115+
private final List<ActivationRule> rules;
116+
117+
public CompositeActivationRule(List<ActivationRule> rules) {
118+
this.rules = new ArrayList<>(rules);
119+
}
120+
121+
@Override
122+
public boolean evaluate(ContextManagerContext context) {
123+
// All rules must evaluate to true (AND logic)
124+
for (ActivationRule rule : rules) {
125+
if (!rule.evaluate(context)) {
126+
return false;
127+
}
128+
}
129+
return true;
130+
}
131+
132+
@Override
133+
public String getDescription() {
134+
StringBuilder sb = new StringBuilder();
135+
sb.append("composite_rule: [");
136+
for (int i = 0; i < rules.size(); i++) {
137+
if (i > 0) {
138+
sb.append(" AND ");
139+
}
140+
sb.append(rules.get(i).getDescription());
141+
}
142+
sb.append("]");
143+
return sb.toString();
144+
}
145+
}
146+
}

0 commit comments

Comments
 (0)