Agent Tracing Storage and Visualization Tutorial

.github/workflows/CI-workflow.yml

@@ -42,7 +42,7 @@ jobs:
     needs: [Get-Require-Approval, Get-CI-Image-Tag, spotless]
     strategy:
       matrix:
-        java: [21, 23]
+        java: [21, 24]
 
     name: Build and Test MLCommons Plugin on linux
     if: github.repository == 'opensearch-project/ml-commons'
@@ -107,7 +107,7 @@ jobs:
     needs: [Get-Require-Approval, Build-ml-linux, spotless]
     strategy:
       matrix:
-        java: [21, 23]
+        java: [21, 24]
 
     name: Test MLCommons Plugin on linux docker
     if: github.repository == 'opensearch-project/ml-commons'
@@ -203,7 +203,7 @@ jobs:
   Build-ml-windows:
     strategy:
       matrix:
-        java: [21, 23]
+        java: [21, 24]
     name: Build and Test MLCommons Plugin on Windows
     if: github.repository == 'opensearch-project/ml-commons'
     needs: [Get-Require-Approval, spotless]

build.gradle

@@ -11,7 +11,7 @@ buildscript {
     ext {
         opensearch_group = "org.opensearch"
         isSnapshot = "true" == System.getProperty("build.snapshot", "true")
-        opensearch_version = System.getProperty("opensearch.version", "3.1.0-SNAPSHOT")
+        opensearch_version = System.getProperty("opensearch.version", "3.2.0-SNAPSHOT")
         buildVersionQualifier = System.getProperty("build.version_qualifier", "")
         asm_version = "9.7"
 
@@ -60,9 +60,9 @@ buildscript {
 }
 
 plugins {
-    id 'com.netflix.nebula.ospackage' version "11.5.0"
+    id 'com.netflix.nebula.ospackage' version "12.0.0"
     id 'java'
-    id "io.freefair.lombok" version "8.4"
+    id "io.freefair.lombok" version "8.14"
     id 'jacoco'
 }
 
@@ -80,7 +80,7 @@ allprojects {
     }
 
     plugins.withId('jacoco') {
-        jacoco.toolVersion = '0.8.12'
+        jacoco.toolVersion = '0.8.13'
     }
 
     project.getExtensions().getExtraProperties().set("versions", VersionProperties.getVersions());

client/src/main/java/org/opensearch/ml/client/MachineLearningNodeClient.java

@@ -214,7 +214,7 @@ public void registerModelGroup(
     @Override
     public void execute(FunctionName name, Input input, ActionListener<MLExecuteTaskResponse> listener) {
         MLExecuteTaskRequest mlExecuteTaskRequest = new MLExecuteTaskRequest(name, input);
-        client.execute(MLExecuteTaskAction.INSTANCE, mlExecuteTaskRequest, listener);
+        client.execute(MLExecuteTaskAction.INSTANCE, mlExecuteTaskRequest, getMLExecuteResponseActionListener(listener));
     }
 
     @Override
@@ -345,6 +345,10 @@ private ActionListener<MLRegisterAgentResponse> getMLRegisterAgentResponseAction
         return wrapActionListener(listener, MLRegisterAgentResponse::fromActionResponse);
     }
 
+    private ActionListener<MLExecuteTaskResponse> getMLExecuteResponseActionListener(ActionListener<MLExecuteTaskResponse> listener) {
+        return wrapActionListener(listener, MLExecuteTaskResponse::fromActionResponse);
+    }
+
     private ActionListener<MLTaskGetResponse> getMLTaskResponseActionListener(ActionListener<MLTask> listener) {
         ActionListener<MLTaskGetResponse> internalListener = ActionListener
             .wrap(getResponse -> { listener.onResponse(getResponse.getMlTask()); }, listener::onFailure);

common/build.gradle

@@ -47,7 +47,7 @@ dependencies {
 }
 
 lombok {
-    version = "1.18.30"
+    version = "1.18.38"
 }
 
 jacocoTestReport {

common/src/main/java/org/opensearch/ml/common/AccessMode.java

@@ -1,8 +1,6 @@
 /*
- *
- *  * Copyright OpenSearch Contributors
- *  * SPDX-License-Identifier: Apache-2.0
- *
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
  */
 
 package org.opensearch.ml.common;

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

@@ -43,6 +43,7 @@ public class CommonValue {
     public static final String ML_AGENT_INDEX = ".plugins-ml-agent";
     public static final String ML_MEMORY_META_INDEX = ".plugins-ml-memory-meta";
     public static final String ML_MEMORY_MESSAGE_INDEX = ".plugins-ml-memory-message";
+    public static final String ML_MEMORY_CONTAINER_INDEX = ".plugins-ml-memory-container";
     public static final String ML_STOP_WORDS_INDEX = ".plugins-ml-stop-words";
     // index used in 2.19 to track MlTaskBatchUpdate task
     public static final String TASK_POLLING_JOB_INDEX = ".ml_commons_task_polling_job";
@@ -63,6 +64,7 @@ public class CommonValue {
     public static final String ML_AGENT_INDEX_MAPPING_PATH = "index-mappings/ml_agent.json";
     public static final String ML_MEMORY_META_INDEX_MAPPING_PATH = "index-mappings/ml_memory_meta.json";
     public static final String ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH = "index-mappings/ml_memory_message.json";
+    public static final String ML_MEMORY_CONTAINER_INDEX_MAPPING_PATH = "index-mappings/ml_memory_container.json";
     public static final String ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_session_management.json";
     public static final String ML_MCP_TOOLS_INDEX_MAPPING_PATH = "index-mappings/ml_mcp_tools.json";
     public static final String ML_JOBS_INDEX_MAPPING_PATH = "index-mappings/ml_jobs.json";

common/src/main/java/org/opensearch/ml/common/FunctionName.java

@@ -31,7 +31,8 @@ public enum FunctionName {
     TEXT_SIMILARITY,
     QUESTION_ANSWERING,
     AGENT,
-    CONNECTOR;
+    CONNECTOR,
+    TOOL;
 
     public static FunctionName from(String value) {
         try {

common/src/main/java/org/opensearch/ml/common/MLCommonsClassLoader.java

@@ -28,6 +28,8 @@
 import org.opensearch.ml.common.output.MLOutputType;
 import org.reflections.Reflections;
 
+import com.fasterxml.jackson.core.JsonParseException;
+
 import lombok.extern.log4j.Log4j2;
 
 @Log4j2
@@ -255,16 +257,18 @@ public static boolean canInitMLInput(FunctionName functionName) {
         return mlInputClassMap.containsKey(functionName);
     }
 
-    public static <S> S initConnector(String name, Object[] initArgs, Class<?>... constructorParameterTypes) {
+    public static <S> S initConnector(String name, Object[] initArgs, Class<?>... constructorParameterTypes) throws JsonParseException {
         return init(connectorClassMap, name, initArgs, constructorParameterTypes);
     }
 
     @SuppressWarnings("unchecked")
-    public static <T extends Enum<T>, S> S initMLInput(T type, Object[] initArgs, Class<?>... constructorParameterTypes) {
+    public static <T extends Enum<T>, S> S initMLInput(T type, Object[] initArgs, Class<?>... constructorParameterTypes)
+        throws JsonParseException {
         return init(mlInputClassMap, type, initArgs, constructorParameterTypes);
     }
 
-    private static <T, S> S init(Map<T, Class<?>> map, T type, Object[] initArgs, Class<?>... constructorParameterTypes) {
+    private static <T, S> S init(Map<T, Class<?>> map, T type, Object[] initArgs, Class<?>... constructorParameterTypes)
+        throws JsonParseException {
         Class<?> clazz = map.get(type);
         if (clazz == null) {
             throw new IllegalArgumentException("Can't find class for type " + type);
@@ -278,6 +282,8 @@ private static <T, S> S init(Map<T, Class<?>> map, T type, Object[] initArgs, Cl
                 throw (MLException) cause;
             } else if (cause instanceof IllegalArgumentException) {
                 throw (IllegalArgumentException) cause;
+            } else if (cause instanceof JsonParseException) {
+                throw (JsonParseException) cause;
             } else {
                 log.error("Failed to init instance for type " + type, e);
                 return null;

common/src/main/java/org/opensearch/ml/common/MLIndex.java

@@ -19,6 +19,8 @@
 import static org.opensearch.ml.common.CommonValue.ML_JOBS_INDEX_MAPPING_PATH;
 import static org.opensearch.ml.common.CommonValue.ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH;
 import static org.opensearch.ml.common.CommonValue.ML_MCP_TOOLS_INDEX_MAPPING_PATH;
+import static org.opensearch.ml.common.CommonValue.ML_MEMORY_CONTAINER_INDEX;
+import static org.opensearch.ml.common.CommonValue.ML_MEMORY_CONTAINER_INDEX_MAPPING_PATH;
 import static org.opensearch.ml.common.CommonValue.ML_MEMORY_MESSAGE_INDEX;
 import static org.opensearch.ml.common.CommonValue.ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH;
 import static org.opensearch.ml.common.CommonValue.ML_MEMORY_META_INDEX;
@@ -45,6 +47,7 @@ public enum MLIndex {
     AGENT(ML_AGENT_INDEX, false, ML_AGENT_INDEX_MAPPING_PATH),
     MEMORY_META(ML_MEMORY_META_INDEX, false, ML_MEMORY_META_INDEX_MAPPING_PATH),
     MEMORY_MESSAGE(ML_MEMORY_MESSAGE_INDEX, false, ML_MEMORY_MESSAGE_INDEX_MAPPING_PATH),
+    MEMORY_CONTAINER(ML_MEMORY_CONTAINER_INDEX, false, ML_MEMORY_CONTAINER_INDEX_MAPPING_PATH),
     MCP_SESSION_MANAGEMENT(MCP_SESSION_MANAGEMENT_INDEX, false, ML_MCP_SESSION_MANAGEMENT_INDEX_MAPPING_PATH),
     MCP_TOOLS(MCP_TOOLS_INDEX, false, ML_MCP_TOOLS_INDEX_MAPPING_PATH),
     JOBS(ML_JOBS_INDEX, false, ML_JOBS_INDEX_MAPPING_PATH);

common/src/main/java/org/opensearch/ml/common/MLModel.java

@@ -965,7 +965,7 @@ Tags getRemoteModelTags(Connector connector) {
      * @param url The URL to analyze for service provider identification
      * @return The identified service provider name, or "unknown" if not found
      */
-    String identifyServiceProvider(String url) {
+    static String identifyServiceProvider(String url) {
         for (String provider : MODEL_SERVICE_PROVIDER_KEYWORDS) {
             if (url.contains(provider)) {
                 return provider;
@@ -975,6 +975,10 @@ String identifyServiceProvider(String url) {
         return TAG_VALUE_UNKNOWN;
     }
 
+    public static String identifyServiceProviderFromUrl(String url) {
+        return identifyServiceProvider(url);
+    }
+
     /**
      * Identifies the model name from the connector configuration using multiple strategies.
      * The method attempts to extract the model name in the following order:

common/src/main/java/org/opensearch/ml/common/MLTask.java

@@ -73,6 +73,7 @@ public class MLTask implements ToXContentObject, Writeable {
     @Setter
     private String error;
     private User user; // TODO: support document level access control later
+    @Setter
     private boolean async;
     @Setter
     private Map<String, Object> remoteJob;

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

@@ -182,7 +182,7 @@ default void validateConnectorURL(List<String> urlRegexes) {
                 }
             }
             if (!hasMatchedUrl) {
-                throw new IllegalArgumentException("Connector URL is not matching the trusted connector endpoint regex, URL is: " + url);
+                throw new IllegalArgumentException("Connector URL is not matching the trusted connector endpoint regex");
             }
         }
     }

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

@@ -9,10 +9,14 @@
 
 import java.io.IOException;
 import java.util.HashSet;
+import java.util.List;
 import java.util.Locale;
 import java.util.Map;
 import java.util.Set;
 
+import org.apache.commons.text.StringSubstitutor;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.common.io.stream.Writeable;
@@ -35,6 +39,17 @@ public class ConnectorAction implements ToXContentObject, Writeable {
     public static final String REQUEST_BODY_FIELD = "request_body";
     public static final String ACTION_PRE_PROCESS_FUNCTION = "pre_process_function";
     public static final String ACTION_POST_PROCESS_FUNCTION = "post_process_function";
+    public static final String OPENAI = "openai";
+    public static final String COHERE = "cohere";
+    public static final String BEDROCK = "bedrock";
+    public static final String SAGEMAKER = "sagemaker";
+    public static final String SAGEMAKER_PRE_POST_FUNC_TEXT = "default";
+    public static final List<String> SUPPORTED_REMOTE_SERVERS_FOR_DEFAULT_ACTION_TYPES = List.of(SAGEMAKER, OPENAI, BEDROCK, COHERE);
+
+    private static final String INBUILT_FUNC_PREFIX = "connector.";
+    private static final String PRE_PROCESS_FUNC = "PreProcessFunction";
+    private static final String POST_PROCESS_FUNC = "PostProcessFunction";
+    private static final Logger logger = LogManager.getLogger(ConnectorAction.class);
 
     private ActionType actionType;
     private String method;
@@ -185,6 +200,81 @@ public static ConnectorAction parse(XContentParser parser) throws IOException {
             .build();
     }
 
+    /**
+     * Checks the compatibility of pre and post-process functions with the selected LLM service.
+     * Each LLM service (eg: Bedrock, OpenAI, SageMaker) has recommended pre and post-process functions
+     * designed for optimal performance. While it's possible to use functions from other services,
+     * it's strongly advised to use the corresponding functions for the best results.
+     * This method logs a warning if non-corresponding functions are detected, but allows the
+     * configuration to proceed. Users should be aware that using mismatched functions may lead
+     * to unexpected behavior or reduced performance, though it won't necessarily cause failures.
+     *
+     * @param parameters - connector parameters
+     */
+    public void validatePrePostProcessFunctions(Map<String, String> parameters) {
+        StringSubstitutor substitutor = new StringSubstitutor(parameters, "${parameters.", "}");
+        String endPoint = substitutor.replace(url);
+        String remoteServer = getRemoteServerFromURL(endPoint);
+        if (!remoteServer.isEmpty()) {
+            validateProcessFunctions(remoteServer, preProcessFunction, PRE_PROCESS_FUNC);
+            validateProcessFunctions(remoteServer, postProcessFunction, POST_PROCESS_FUNC);
+        }
+    }
+
+    /**
+     * To get the remote server name from url
+     *
+     * @param url - remote server url
+     * @return - returns the corresponding remote server name for url, if server is not in the pre-defined list,
+     * it returns null
+     */
+    public static String getRemoteServerFromURL(String url) {
+        return SUPPORTED_REMOTE_SERVERS_FOR_DEFAULT_ACTION_TYPES.stream().filter(url::contains).findFirst().orElse("");
+    }
+
+    private void validateProcessFunctions(String remoteServer, String processFunction, String funcNameForWarnText) {
+        if (isInBuiltProcessFunction(processFunction)) {
+            switch (remoteServer) {
+                case OPENAI:
+                    if (!processFunction.contains(OPENAI)) {
+                        logWarningForInvalidProcessFunc(OPENAI, funcNameForWarnText);
+                    }
+                    break;
+                case COHERE:
+                    if (!processFunction.contains(COHERE)) {
+                        logWarningForInvalidProcessFunc(COHERE, funcNameForWarnText);
+                    }
+                    break;
+                case BEDROCK:
+                    if (!processFunction.contains(BEDROCK)) {
+                        logWarningForInvalidProcessFunc(BEDROCK, funcNameForWarnText);
+                    }
+                    break;
+                case SAGEMAKER:
+                    if (!processFunction.contains(SAGEMAKER_PRE_POST_FUNC_TEXT)) {
+                        logWarningForInvalidProcessFunc(SAGEMAKER, funcNameForWarnText);
+                    }
+            }
+        }
+    }
+
+    private boolean isInBuiltProcessFunction(String processFunction) {
+        return (processFunction != null && processFunction.startsWith(INBUILT_FUNC_PREFIX));
+    }
+
+    private void logWarningForInvalidProcessFunc(String remoteServer, String funcNameForWarnText) {
+        logger
+            .warn(
+                "LLM service is "
+                    + remoteServer
+                    + ", so "
+                    + funcNameForWarnText
+                    + " should be corresponding to "
+                    + remoteServer
+                    + " for better results."
+            );
+    }
+
     public enum ActionType {
         PREDICT,
         EXECUTE,

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

@@ -70,6 +70,11 @@ public HttpConnector(
         String tenantId
     ) {
         validateProtocol(protocol);
+        if (actions != null) {
+            for (ConnectorAction action : actions) {
+                action.validatePrePostProcessFunctions(parameters);
+            }
+        }
         this.name = name;
         this.description = description;
         this.version = version;

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

@@ -438,7 +438,7 @@ public void validateConnectorURL(List<String> urlRegexes) {
             }
         }
         if (!hasMatchedUrl) {
-            throw new IllegalArgumentException("Connector URL is not matching the trusted connector endpoint regex, URL is: " + url);
+            throw new IllegalArgumentException("Connector URL is not matching the trusted connector endpoint regex");
         }
     }
 

common/src/main/java/org/opensearch/ml/common/connector/functions/preprocess/MultiModalConnectorPreProcessFunction.java

@@ -1,8 +1,6 @@
 /*
- *
- *  * Copyright OpenSearch Contributors
- *  * SPDX-License-Identifier: Apache-2.0
- *
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
  */
 
 package org.opensearch.ml.common.connector.functions.preprocess;