Prevent qp tool create when not enabled

Signed-off-by: rithin-pullela-aws <rithinp@amazon.com>

Author: rithin-pullela-aws

ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/QueryPlanningTool.java

@@ -5,11 +5,15 @@
 
 package org.opensearch.ml.engine.tools;
 
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENTIC_SEARCH_DISABLED_MESSAGE;
+
 import java.util.List;
 import java.util.Map;
 
 import org.apache.commons.text.StringSubstitutor;
+import org.opensearch.OpenSearchException;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.common.spi.tools.ToolAnnotation;
 import org.opensearch.ml.common.spi.tools.WithModelTool;
 import org.opensearch.ml.repackage.com.google.common.annotations.VisibleForTesting;
@@ -108,6 +112,7 @@ public boolean validate(Map<String, String> parameters) {
     public static class Factory implements WithModelTool.Factory<QueryPlanningTool> {
         private Client client;
         private static volatile Factory INSTANCE;
+        private static MLFeatureEnabledSetting mlFeatureEnabledSetting;
 
         public static Factory getInstance() {
             if (INSTANCE != null) {
@@ -122,13 +127,18 @@ public static Factory getInstance() {
             }
         }
 
-        public void init(Client client) {
+        public void init(Client client, MLFeatureEnabledSetting mlFeatureEnabledSetting) {
             this.client = client;
+            this.mlFeatureEnabledSetting = mlFeatureEnabledSetting;
         }
 
         @Override
         public QueryPlanningTool create(Map<String, Object> map) {
 
+            if (!mlFeatureEnabledSetting.isAgenticSearchEnabled()) {
+                throw new OpenSearchException(ML_COMMONS_AGENTIC_SEARCH_DISABLED_MESSAGE);
+            }
+
             MLModelTool queryGenerationTool = MLModelTool.Factory.getInstance().create(map);
 
             String type = (String) map.get(GENERATION_TYPE_FIELD);

ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/QueryPlanningToolTests.java

@@ -14,6 +14,8 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import static org.opensearch.ml.common.settings.MLCommonsSettings.ML_COMMONS_AGENTIC_SEARCH_DISABLED_MESSAGE;
 import static org.opensearch.ml.engine.tools.QueryPlanningTool.DEFAULT_DESCRIPTION;
 import static org.opensearch.ml.engine.tools.QueryPlanningTool.MODEL_ID_FIELD;
 
@@ -31,7 +33,9 @@
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.MockitoAnnotations;
+import org.opensearch.OpenSearchException;
 import org.opensearch.core.action.ActionListener;
+import org.opensearch.ml.common.settings.MLFeatureEnabledSetting;
 import org.opensearch.ml.common.spi.tools.Tool;
 import org.opensearch.transport.client.Client;
 
@@ -46,6 +50,9 @@ public class QueryPlanningToolTests {
     @Mock
     private MLModelTool queryGenerationTool;
 
+    @Mock
+    private MLFeatureEnabledSetting mlFeatureEnabledSetting;
+
     private Map<String, String> validParams;
     private Map<String, String> emptyParams;
 
@@ -55,7 +62,14 @@ public class QueryPlanningToolTests {
     public void setup() {
         MockitoAnnotations.openMocks(this);
         MLModelTool.Factory.getInstance().init(client);
-        factory = new QueryPlanningTool.Factory();
+
+        // Mock the MLFeatureEnabledSetting to return true for agentic search
+        when(mlFeatureEnabledSetting.isAgenticSearchEnabled()).thenReturn(true);
+
+        // Initialize the factory with mocked dependencies
+        factory = QueryPlanningTool.Factory.getInstance();
+        factory.init(client, mlFeatureEnabledSetting);
+
         validParams = new HashMap<>();
         validParams.put("prompt", "test prompt");
         emptyParams = Collections.emptyMap();
@@ -269,4 +283,16 @@ public void testFactoryCreateWithInvalidType() {
         Exception exception = assertThrows(IllegalArgumentException.class, () -> factory.create(map));
         assertEquals("Invalid generation type: invalid. The current supported types are llmGenerated.", exception.getMessage());
     }
+
+    @Test
+    public void testFactoryCreateWhenAgenticSearchDisabled() {
+        // Mock the MLFeatureEnabledSetting to return false for agentic search
+        when(mlFeatureEnabledSetting.isAgenticSearchEnabled()).thenReturn(false);
+        
+        Map<String, Object> map = new HashMap<>();
+        map.put(QueryPlanningTool.MODEL_ID_FIELD, "modelId");
+
+        Exception exception = assertThrows(OpenSearchException.class, () -> factory.create(map));
+        assertEquals(ML_COMMONS_AGENTIC_SEARCH_DISABLED_MESSAGE, exception.getMessage());
+    }
 }

plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java

@@ -730,7 +730,7 @@ public Collection<Object> createComponents(
         SearchIndexTool.Factory.getInstance().init(client, xContentRegistry);
         VisualizationsTool.Factory.getInstance().init(client);
         ConnectorTool.Factory.getInstance().init(client);
-        QueryPlanningTool.Factory.getInstance().init(client);
+        QueryPlanningTool.Factory.getInstance().init(client, mlFeatureEnabledSetting);
 
         toolFactories.put(MLModelTool.TYPE, MLModelTool.Factory.getInstance());
         toolFactories.put(McpSseTool.TYPE, McpSseTool.Factory.getInstance());
@@ -1179,7 +1179,9 @@ public List<Setting<?>> getSettings() {
                 MLCommonsSettings.ML_COMMONS_MCP_CONNECTOR_ENABLED,
                 MLCommonsSettings.ML_COMMONS_MCP_SERVER_ENABLED,
                 MLCommonsSettings.ML_COMMONS_METRIC_COLLECTION_ENABLED,
-                MLCommonsSettings.ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED
+                MLCommonsSettings.ML_COMMONS_STATIC_METRIC_COLLECTION_ENABLED,
+                MLCommonsSettings.ML_COMMONS_EXECUTE_TOOL_ENABLED,
+                MLCommonsSettings.ML_COMMONS_AGENTIC_SEARCH_ENABLED
             );
         return settings;
     }