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]>

Author: mingshl

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

@@ -365,10 +365,6 @@ public <T> T createPayload(String action, Map<String, String> parameters) {
                 jsonObject.addProperty("stream", true);
                 payload = jsonObject.toString();
             }
-            // Log payload for debugging
-
-            log.info("=== PAYLOAD DEBUG === Action: {} | Payload: {}", action, payload);
-
             return (T) payload;
         }
         return (T) parameters.get("http_body");

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

@@ -65,7 +65,7 @@ public class ContextManagerContext {
      * Additional parameters for context processing
      */
     @Builder.Default
-    private Map<String, Object> parameters = new HashMap<>();
+    private Map<String, String> parameters = new HashMap<>();
 
     /**
      * Get the total token count for the current context.
@@ -174,7 +174,7 @@ public Object getParameter(String key) {
      * @param key the parameter key
      * @param value the parameter value
      */
-    public void setParameter(String key, Object value) {
+    public void setParameter(String key, String value) {
         if (parameters == null) {
             parameters = new HashMap<>();
         }

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

@@ -110,7 +110,12 @@ public AgentMLInput(XContentParser parser, FunctionName functionName) throws IOE
                     tenantId = parser.textOrNull();
                     break;
                 case PARAMETERS_FIELD:
-                    Map<String, String> parameters = StringUtils.getParameterMap(parser.map());
+                    Map<String, Object> parameterObjs = parser.map();
+                    Map<String, String> parameters = StringUtils.getParameterMap(parameterObjs);
+                    // Extract context_management from parameters
+                    if (parameterObjs.containsKey("context_management")) {
+                        contextManagementName = (String) parameterObjs.get("context_management");
+                    }
                     inputDataset = new RemoteInferenceInputDataSet(parameters);
                     break;
                 case ASYNC_FIELD: