add multi tenacy support

Signed-off-by: Hailong Cui <[email protected]>
opensearch-project · Jan 24, 2025 · 5144fdf · 5144fdf
1 parent 76f1cfc
commit 5144fdf
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 3 deletions.
diff --git a/src/main/java/org/opensearch/agent/tools/CreateAnomalyDetectorTool.java b/src/main/java/org/opensearch/agent/tools/CreateAnomalyDetectorTool.java
@@ -5,6 +5,7 @@
 
 package org.opensearch.agent.tools;
 
+import static org.opensearch.ml.common.CommonValue.TENANT_ID_FIELD;
 import static org.opensearch.ml.common.utils.StringUtils.gson;
 
 import java.io.IOException;
@@ -165,6 +166,7 @@ public CreateAnomalyDetectorTool(Client client, String modelId, String modelType
      */
     @Override
     public <T> void run(Map<String, String> parameters, ActionListener<T> listener) {
+        final String tenantId = parameters.get(TENANT_ID_FIELD);
         Map<String, String> enrichedParameters = enrichParameters(parameters);
         String indexName = enrichedParameters.get("index");
         if (Strings.isNullOrEmpty(indexName)) {
@@ -228,7 +230,8 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
             ActionRequest request = new MLPredictionTaskRequest(
                 modelId,
                 MLInput.builder().algorithm(FunctionName.REMOTE).inputDataset(inputDataSet).build(),
-                null
+                null,
+                tenantId
             );
 
             client.execute(MLPredictionTaskAction.INSTANCE, request, ActionListener.wrap(mlTaskResponse -> {

diff --git a/src/main/java/org/opensearch/agent/tools/PPLTool.java b/src/main/java/org/opensearch/agent/tools/PPLTool.java
@@ -5,6 +5,8 @@
 
 package org.opensearch.agent.tools;
 
+import static org.opensearch.ml.common.CommonValue.TENANT_ID_FIELD;
+
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
@@ -172,6 +174,7 @@ public PPLTool(
     @SuppressWarnings("unchecked")
     @Override
     public <T> void run(Map<String, String> parameters, ActionListener<T> listener) {
+        final String tenantId = parameters.get(TENANT_ID_FIELD);
         extractFromChatParameters(parameters);
         String indexName = getIndexNameFromParameters(parameters);
         if (StringUtils.isBlank(indexName)) {
@@ -209,7 +212,8 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
                 ActionRequest request = new MLPredictionTaskRequest(
                     modelId,
                     MLInput.builder().algorithm(FunctionName.REMOTE).inputDataset(inputDataSet).build(),
-                    null
+                    null,
+                    tenantId
                 );
                 client.execute(MLPredictionTaskAction.INSTANCE, request, ActionListener.wrap(mlTaskResponse -> {
                     ModelTensorOutput modelTensorOutput = (ModelTensorOutput) mlTaskResponse.getOutput();

diff --git a/src/main/java/org/opensearch/agent/tools/RAGTool.java b/src/main/java/org/opensearch/agent/tools/RAGTool.java
@@ -7,6 +7,7 @@
 
 import static org.apache.commons.lang3.StringEscapeUtils.escapeJson;
 import static org.opensearch.agent.tools.AbstractRetrieverTool.*;
+import static org.opensearch.ml.common.CommonValue.TENANT_ID_FIELD;
 import static org.opensearch.ml.common.utils.StringUtils.gson;
 import static org.opensearch.ml.common.utils.StringUtils.toJson;
 
@@ -94,6 +95,8 @@ public Object parse(Object o) {
     }
 
     public <T> void run(Map<String, String> parameters, ActionListener<T> listener) {
+        final String tenantId = parameters.get(TENANT_ID_FIELD);
+
         String input = null;
 
         if (!this.validate(parameters)) {
@@ -145,7 +148,7 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
 
             RemoteInferenceInputDataSet inputDataSet = RemoteInferenceInputDataSet.builder().parameters(tmpParameters).build();
             MLInput mlInput = MLInput.builder().algorithm(FunctionName.REMOTE).inputDataset(inputDataSet).build();
-            ActionRequest request = new MLPredictionTaskRequest(this.inferenceModelId, mlInput, null);
+            ActionRequest request = new MLPredictionTaskRequest(this.inferenceModelId, mlInput, null, tenantId);
 
             client.execute(MLPredictionTaskAction.INSTANCE, request, ActionListener.wrap(resp -> {
                 ModelTensorOutput modelTensorOutput = (ModelTensorOutput) resp.getOutput();