update guardrails

Signed-off-by: Jing Zhang <[email protected]>

Author: jngz-es

common/src/main/java/org/opensearch/ml/common/transport/model/MLUpdateModelInput.java

@@ -15,6 +15,7 @@
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentParser;
 import org.opensearch.ml.common.connector.Connector;
+import org.opensearch.ml.common.model.Guardrails;
 import org.opensearch.ml.common.model.MLModelConfig;
 import org.opensearch.ml.common.controller.MLRateLimiter;
 import org.opensearch.ml.common.model.TextEmbeddingModelConfig;
@@ -43,6 +44,7 @@ public class MLUpdateModelInput implements ToXContentObject, Writeable {
     public static final String CONNECTOR_FIELD = "connector"; // optional
     public static final String LAST_UPDATED_TIME_FIELD = "last_updated_time"; // passively set when sending update
                                                                               // request
+    public static final String GUARDRAILS_FIELD = "guardrails";
 
     @Getter
     private String modelId;
@@ -57,11 +59,12 @@ public class MLUpdateModelInput implements ToXContentObject, Writeable {
     private String connectorId;
     private MLCreateConnectorInput connector;
     private Instant lastUpdateTime;
+    private Guardrails guardrails;
 
     @Builder(toBuilder = true)
     public MLUpdateModelInput(String modelId, String description, String version, String name, String modelGroupId,
             Boolean isEnabled, MLRateLimiter rateLimiter, MLModelConfig modelConfig,
-            Connector updatedConnector, String connectorId, MLCreateConnectorInput connector, Instant lastUpdateTime) {
+            Connector updatedConnector, String connectorId, MLCreateConnectorInput connector, Instant lastUpdateTime, Guardrails guardrails) {
         this.modelId = modelId;
         this.description = description;
         this.version = version;
@@ -74,6 +77,7 @@ public MLUpdateModelInput(String modelId, String description, String version, St
         this.connectorId = connectorId;
         this.connector = connector;
         this.lastUpdateTime = lastUpdateTime;
+        this.guardrails = guardrails;
     }
 
     public MLUpdateModelInput(StreamInput in) throws IOException {
@@ -97,6 +101,9 @@ public MLUpdateModelInput(StreamInput in) throws IOException {
             connector = new MLCreateConnectorInput(in);
         }
         lastUpdateTime = in.readOptionalInstant();
+        if (in.readBoolean()) {
+            this.guardrails = new Guardrails(in);
+        }
     }
 
     @Override
@@ -136,6 +143,9 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         if (lastUpdateTime != null) {
             builder.field(LAST_UPDATED_TIME_FIELD, lastUpdateTime.toEpochMilli());
         }
+        if (guardrails != null) {
+            builder.field(GUARDRAILS_FIELD, guardrails);
+        }
         builder.endObject();
         return builder;
     }
@@ -174,6 +184,9 @@ public XContentBuilder toXContentForUpdateRequestDoc(XContentBuilder builder, Pa
         if (lastUpdateTime != null) {
             builder.field(LAST_UPDATED_TIME_FIELD, lastUpdateTime.toEpochMilli());
         }
+        if (guardrails != null) {
+            builder.field(GUARDRAILS_FIELD, guardrails);
+        }
         builder.endObject();
         return builder;
     }
@@ -212,6 +225,12 @@ public void writeTo(StreamOutput out) throws IOException {
             out.writeBoolean(false);
         }
         out.writeOptionalInstant(lastUpdateTime);
+        if (guardrails != null) {
+            out.writeBoolean(true);
+            guardrails.writeTo(out);
+        } else {
+            out.writeBoolean(false);
+        }
     }
 
     public static MLUpdateModelInput parse(XContentParser parser) throws IOException {
@@ -227,6 +246,7 @@ public static MLUpdateModelInput parse(XContentParser parser) throws IOException
         String connectorId = null;
         MLCreateConnectorInput connector = null;
         Instant lastUpdateTime = null;
+        Guardrails guardrails = null;
 
         ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.currentToken(), parser);
         while (parser.nextToken() != XContentParser.Token.END_OBJECT) {
@@ -257,6 +277,9 @@ public static MLUpdateModelInput parse(XContentParser parser) throws IOException
                 case CONNECTOR_FIELD:
                     connector = MLCreateConnectorInput.parse(parser, true);
                     break;
+                case GUARDRAILS_FIELD:
+                    guardrails = Guardrails.parse(parser);
+                    break;
                 default:
                     parser.skipChildren();
                     break;
@@ -265,6 +288,6 @@ public static MLUpdateModelInput parse(XContentParser parser) throws IOException
         // Model ID can only be set through RestRequest. Model version can only be set
         // automatically.
         return new MLUpdateModelInput(modelId, description, version, name, modelGroupId, isEnabled, rateLimiter,
-                modelConfig, updatedConnector, connectorId, connector, lastUpdateTime);
+                modelConfig, updatedConnector, connectorId, connector, lastUpdateTime, guardrails);
     }
 }

plugin/src/main/java/org/opensearch/ml/action/models/UpdateModelTransportAction.java

@@ -206,11 +206,13 @@ private void updateRemoteOrTextEmbeddingModel(
         String newConnectorId = Strings.hasLength(updateModelInput.getConnectorId()) ? updateModelInput.getConnectorId() : null;
         boolean isModelDeployed = isModelDeployed(mlModel.getModelState());
         // This flag is used to decide if we need to re-deploy the predictor(model) when updating the model cache.
-        // If one of the internal connector, stand-alone connector id, model quota flag, as well as the model rate limiter needs update, we
+        // If one of the internal connector, stand-alone connector id, model quota flag, as well as the model rate limiter and guardrails
+        // need update, we
         // need to perform a re-deploy.
         boolean isPredictorUpdate = (updateModelInput.getConnector() != null)
             || (newConnectorId != null)
-            || !Objects.equals(updateModelInput.getIsEnabled(), mlModel.getIsEnabled());
+            || !Objects.equals(updateModelInput.getIsEnabled(), mlModel.getIsEnabled())
+            || (updateModelInput.getGuardrails() != null);
         if (MLRateLimiter.updateValidityPreCheck(mlModel.getRateLimiter(), updateModelInput.getRateLimiter())) {
             MLRateLimiter updatedRateLimiterConfig = MLRateLimiter.update(mlModel.getRateLimiter(), updateModelInput.getRateLimiter());
             updateModelInput.setRateLimiter(updatedRateLimiterConfig);