Add remote guardrails model support (#7377)

* remote guardrails model support Signed-off-by: Jing Zhang <[email protected]> * address comments Signed-off-by: Jing Zhang <[email protected]> * Add guardrail information Signed-off-by: Fanit Kolchina <[email protected]> * Rename response accept parameter Signed-off-by: Fanit Kolchina <[email protected]> * typo fix Signed-off-by: Fanit Kolchina <[email protected]> * Fix links Signed-off-by: Fanit Kolchina <[email protected]> * Apply suggestions from code review Co-authored-by: Nathan Bower <[email protected]> Signed-off-by: kolchfa-aws <[email protected]> * Update _ml-commons-plugin/remote-models/guardrails.md Signed-off-by: kolchfa-aws <[email protected]> --------- Signed-off-by: Jing Zhang <[email protected]> Signed-off-by: Fanit Kolchina <[email protected]> Signed-off-by: kolchfa-aws <[email protected]> Co-authored-by: Fanit Kolchina <[email protected]> Co-authored-by: kolchfa-aws <[email protected]> Co-authored-by: Nathan Bower <[email protected]>
opensearch-project · Jun 21, 2024 · 8f71179 · 8f71179
1 parent a4437a0
commit 8f71179
Show file tree

Hide file tree

Showing 3 changed files with 417 additions and 19 deletions.
diff --git a/_ml-commons-plugin/api/model-apis/register-model.md b/_ml-commons-plugin/api/model-apis/register-model.md
@@ -260,15 +260,24 @@ To register an externally hosted model with guardrails, provide the `guardrails`
 
 Field | Data type | Description
 :---  | :--- | :---
-`type` | String | The guardrail type. Currently, only `local_regex` is supported.
-`input_guardrail`| Object |  The guardrail for the model input. |
-`output_guardrail`| Object |  The guardrail for the model output. |
-`stop_words`| Object | The list of indexes containing stopwords used for the model input/output validation. If the model prompt/response contains a stopword contained in any of the indexes, the predict request on this model is rejected. |
-`index_name`| Object | The name of the index storing the stopwords. |
-`source_fields`| Object | The name of the field storing the stopwords. |
-`regex`| Object |  A regular expression used for input/output validation. If the model prompt/response matches the regular expression, the predict request on this model is rejected. |
+`type` | String | The guardrail type. Valid values are [`local_regex`](#example-request-regex-and-stopword-validation) and [`model`](#example-request-guardrail-model-validation). Using `local_regex`, you can specify a regular expression or stop words. Using `model`, you can specify a guardrail model. For more information, see [Guardrails]({{site.url}}{{site.baseurl}}/ml-commons-plugin/remote-models/guardrails/). 
+`input_guardrail`| Object |  The guardrail for the model input. 
+`output_guardrail`| Object |  The guardrail for the model output. 
+`stop_words`| Object | The list of indexes containing stopwords used for model input/output validation. If the model prompt/response contains a stopword contained in any of the indexes, then the predict request on the model is rejected. 
+`index_name`| Object | The name of the index storing the stopwords. 
+`source_fields`| Object | The name of the field storing the stopwords. 
+`regex`| Object |  A regular expression used for input/output validation. If the model prompt/response matches the regular expression, then the predict request on the model is rejected. 
+`model_id`| String  | The guardrail model used to validate user input and LLM output. 
+`response_filter`| String | The dot path of the field containing the guardrail model response. 
+`response_validation_regex`| String | The regular expression used to validate the guardrail model response.     
 
-#### Example request: Externally hosted model with guardrails
+## Examples
+
+The following examples configure an externally hosted model with guardrails.
+
+#### Example request: Regex and stopword validation
+
+The following example uses a regular expression and a set of stopwords to validate the LLM response:
 
 ```json
 POST /_plugins/_ml/models/_register
@@ -303,7 +312,36 @@ POST /_plugins/_ml/models/_register
 ```
 {% include copy-curl.html %}
 
-For a complete example, see [Guardrails]({{site.url}}{{site.baseurl}}/ml-commons-plugin/remote-models/guardrails/).
+For a complete example, see [Validating input/output using stopwords and regex]({{site.url}}{{site.baseurl}}/ml-commons-plugin/remote-models/guardrails/#validating-inputoutput-using-stopwords-and-regex).
+
+#### Example request: Guardrail model validation
+
+The following example uses a guardrail model to validate the LLM response:
+
+```json
+POST /_plugins/_ml/models/_register?deploy=true
+{
+    "name": "Bedrock Claude V2 model with guardrails model",
+    "function_name": "remote",
+    "model_group_id": "ppSmpo8Bi-GZ0tf1i7cD",
+    "description": "Bedrock Claude V2 model with guardrails model",
+    "connector_id": "xnJjDZABNFJeYR3IPvTO",
+    "guardrails": {
+        "input_guardrail": {
+            "model_id": "o3JaDZABNFJeYR3I2fRV",
+            "response_validation_regex": "^\\s*\"[Aa]ccept\"\\s*$"
+        },
+        "output_guardrail": {
+            "model_id": "o3JaDZABNFJeYR3I2fRV",
+            "response_validation_regex": "^\\s*\"[Aa]ccept\"\\s*$"
+        },
+        "type": "model"
+    }
+}
+```
+{% include copy-curl.html %}
+
+For a complete example, see [Validating input/output using a guardrail model]({{site.url}}{{site.baseurl}}/ml-commons-plugin/remote-models/guardrails/#validating-inputoutput-using-a-guardrail-model).
 
 #### Example response
 

diff --git a/_ml-commons-plugin/api/model-apis/update-model.md b/_ml-commons-plugin/api/model-apis/update-model.md
@@ -64,12 +64,13 @@ PUT /_plugins/_ml/models/T_S-cY0BKCJ3ot9qr0aP
 ```
 {% include copy-curl.html %}
 
-#### Example request: Updating the guardrails
+#### Example requests: Updating the guardrails
 
 ```json
 PUT /_plugins/_ml/models/MzcIJX8BA7mbufL6DOwl
 {
   "guardrails": {
+    "type": "local_regex",
     "input_guardrail": {
       "stop_words": [
         {
@@ -93,6 +94,24 @@ PUT /_plugins/_ml/models/MzcIJX8BA7mbufL6DOwl
 ```
 {% include copy-curl.html %}
 
+```json
+PUT /_plugins/_ml/models/9uGdCJABjaMXYrp14YRj
+{
+  "guardrails": {
+    "type": "model",
+    "input_guardrail": {
+      "model_id": "V-G1CJABjaMXYrp1QoUC",
+      "response_validation_regex": "^\\s*[Aa]ccept\\s*$"
+    },
+    "output_guardrail": {
+      "model_id": "V-G1CJABjaMXYrp1QoUC",
+      "response_validation_regex": "^\\s*[Aa]ccept\\s*$"
+    }
+  }
+}
+```
+{% include copy-curl.html %}
+
 #### Example response
 
 ```json