Skip to content

Commit f6d461d

Browse files
committed
Add Context Manager to PER (opensearch-project#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]>
1 parent 483e8b9 commit f6d461d

File tree

15 files changed

+285
-563
lines changed

15 files changed

+285
-563
lines changed

common/src/main/java/org/opensearch/ml/common/connector/HttpConnector.java

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -369,10 +369,6 @@ public <T> T createPayload(String action, Map<String, String> parameters) {
369369
payload = jsonObject.toString();
370370
}
371371
}
372-
// Log payload for debugging
373-
374-
log.info("=== PAYLOAD DEBUG === Action: {} | Payload: {}", action, payload);
375-
376372
return (T) payload;
377373
}
378374
return (T) parameters.get("http_body");

common/src/main/java/org/opensearch/ml/common/contextmanager/ContextManagerContext.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -65,7 +65,7 @@ public class ContextManagerContext {
6565
* Additional parameters for context processing
6666
*/
6767
@Builder.Default
68-
private Map<String, Object> parameters = new HashMap<>();
68+
private Map<String, String> parameters = new HashMap<>();
6969

7070
/**
7171
* Get the total token count for the current context.
@@ -174,7 +174,7 @@ public Object getParameter(String key) {
174174
* @param key the parameter key
175175
* @param value the parameter value
176176
*/
177-
public void setParameter(String key, Object value) {
177+
public void setParameter(String key, String value) {
178178
if (parameters == null) {
179179
parameters = new HashMap<>();
180180
}

common/src/main/java/org/opensearch/ml/common/input/execute/agent/AgentMLInput.java

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -149,7 +149,12 @@ public AgentMLInput(XContentParser parser, FunctionName functionName) throws IOE
149149
break;
150150
case PARAMETERS_FIELD:
151151
// Legacy format - parse parameters into RemoteInferenceInputDataSet
152-
Map<String, String> parameters = StringUtils.getParameterMap(parser.map());
152+
Map<String, Object> parameterObjs = parser.map();
153+
Map<String, String> parameters = StringUtils.getParameterMap(parameterObjs);
154+
// Extract context_management from parameters
155+
if (parameterObjs.containsKey("context_management")) {
156+
contextManagementName = (String) parameterObjs.get("context_management");
157+
}
153158
inputDataset = new RemoteInferenceInputDataSet(parameters);
154159
break;
155160
case INPUT_FIELD:

common/src/test/java/org/opensearch/ml/common/contextmanager/ToolsOutputTruncateManagerTest.java

Lines changed: 0 additions & 266 deletions
This file was deleted.

0 commit comments

Comments
 (0)