[Backport 2.13] Added new Guardrail field for remote model (#624)

Added new Guardrail field for remote model (#622)

* Added new field guarddail for remote model



* Fixed parsing



* Deserialize



* fixing guardrails



* Added break



---------




(cherry picked from commit 4a12730)

Signed-off-by: Owais Kazi <[email protected]>
Signed-off-by: Joshua Palis <[email protected]>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Joshua Palis <[email protected]>

Author: 3 people

src/main/java/org/opensearch/flowframework/common/CommonValue.java

@@ -168,6 +168,8 @@ private CommonValue() {}
     public static final String PIPELINE_ID = "pipeline_id";
     /** Pipeline Configurations */
     public static final String CONFIGURATIONS = "configurations";
+    /** Guardrails field */
+    public static final String GUARDRAILS_FIELD = "guardrails";
 
     /*
      * Constants associated with resource provisioning / state

src/main/java/org/opensearch/flowframework/model/WorkflowNode.java

@@ -20,6 +20,7 @@
 import org.opensearch.flowframework.workflow.ProcessNode;
 import org.opensearch.flowframework.workflow.WorkflowData;
 import org.opensearch.flowframework.workflow.WorkflowStep;
+import org.opensearch.ml.common.model.Guardrails;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -32,6 +33,7 @@
 import static java.util.concurrent.TimeUnit.SECONDS;
 import static org.opensearch.core.xcontent.XContentParserUtils.ensureExpectedToken;
 import static org.opensearch.flowframework.common.CommonValue.CONFIGURATIONS;
+import static org.opensearch.flowframework.common.CommonValue.GUARDRAILS_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.TOOLS_ORDER_FIELD;
 import static org.opensearch.flowframework.util.ParseUtils.buildStringToObjectMap;
 import static org.opensearch.flowframework.util.ParseUtils.buildStringToStringMap;
@@ -95,6 +97,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
             xContentBuilder.field(e.getKey());
             if (e.getValue() instanceof String || e.getValue() instanceof Number || e.getValue() instanceof Boolean) {
                 xContentBuilder.value(e.getValue());
+            } else if (GUARDRAILS_FIELD.equals(e.getKey())) {
+                Guardrails g = (Guardrails) e.getValue();
+                xContentBuilder.value(g);
             } else if (e.getValue() instanceof Map<?, ?>) {
                 buildStringToStringMap(xContentBuilder, (Map<?, ?>) e.getValue());
             } else if (e.getValue() instanceof Object[]) {
@@ -156,13 +161,16 @@ public static WorkflowNode parse(XContentParser parser) throws IOException {
                                 userInputs.put(inputFieldName, parser.text());
                                 break;
                             case START_OBJECT:
-                                if (CONFIGURATIONS.equals(inputFieldName)) {
+                                if (GUARDRAILS_FIELD.equals(inputFieldName)) {
+                                    userInputs.put(inputFieldName, Guardrails.parse(parser));
+                                    break;
+                                } else if (CONFIGURATIONS.equals(inputFieldName)) {
                                     Map<String, Object> configurationsMap = parser.map();
                                     try {
                                         String configurationsString = ParseUtils.parseArbitraryStringToObjectMapToString(configurationsMap);
                                         userInputs.put(inputFieldName, configurationsString);
                                     } catch (Exception ex) {
-                                        String errorMessage = "Failed to parse configuration map";
+                                        String errorMessage = "Failed to parse" + inputFieldName + "map";
                                         logger.error(errorMessage, ex);
                                         throw new FlowFrameworkException(errorMessage, RestStatus.BAD_REQUEST);
                                     }

src/main/java/org/opensearch/flowframework/workflow/RegisterRemoteModelStep.java

@@ -20,6 +20,7 @@
 import org.opensearch.flowframework.util.ParseUtils;
 import org.opensearch.ml.client.MachineLearningNodeClient;
 import org.opensearch.ml.common.FunctionName;
+import org.opensearch.ml.common.model.Guardrails;
 import org.opensearch.ml.common.transport.register.MLRegisterModelInput;
 import org.opensearch.ml.common.transport.register.MLRegisterModelInput.MLRegisterModelInputBuilder;
 import org.opensearch.ml.common.transport.register.MLRegisterModelResponse;
@@ -29,6 +30,7 @@
 
 import static org.opensearch.flowframework.common.CommonValue.DEPLOY_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.DESCRIPTION_FIELD;
+import static org.opensearch.flowframework.common.CommonValue.GUARDRAILS_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.NAME_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.REGISTER_MODEL_STATUS;
 import static org.opensearch.flowframework.common.WorkflowResources.CONNECTOR_ID;
@@ -71,7 +73,7 @@ public PlainActionFuture<WorkflowData> execute(
         PlainActionFuture<WorkflowData> registerRemoteModelFuture = PlainActionFuture.newFuture();
 
         Set<String> requiredKeys = Set.of(NAME_FIELD, CONNECTOR_ID);
-        Set<String> optionalKeys = Set.of(MODEL_GROUP_ID, DESCRIPTION_FIELD, DEPLOY_FIELD);
+        Set<String> optionalKeys = Set.of(MODEL_GROUP_ID, DESCRIPTION_FIELD, DEPLOY_FIELD, GUARDRAILS_FIELD);
 
         try {
             Map<String, Object> inputs = ParseUtils.getInputsFromPreviousSteps(
@@ -87,6 +89,7 @@ public PlainActionFuture<WorkflowData> execute(
             String modelGroupId = (String) inputs.get(MODEL_GROUP_ID);
             String description = (String) inputs.get(DESCRIPTION_FIELD);
             String connectorId = (String) inputs.get(CONNECTOR_ID);
+            Guardrails guardRails = (Guardrails) inputs.get(GUARDRAILS_FIELD);
             final Boolean deploy = (Boolean) inputs.get(DEPLOY_FIELD);
 
             MLRegisterModelInputBuilder builder = MLRegisterModelInput.builder()
@@ -103,6 +106,11 @@ public PlainActionFuture<WorkflowData> execute(
             if (deploy != null) {
                 builder.deployModel(deploy);
             }
+
+            if (guardRails != null) {
+                builder.guardrails(guardRails);
+            }
+
             MLRegisterModelInput mlInput = builder.build();
 
             mlClient.register(mlInput, new ActionListener<MLRegisterModelResponse>() {