Merge branch 'main' into agent_bug_fix

Author: dhrubo-os

common/src/main/java/org/opensearch/ml/common/transport/connector/MLCreateConnectorRequest.java

@@ -46,7 +46,9 @@ public ActionRequestValidationException validate() {
             return addValidationError("ML Connector input can't be null", null);
         }
         Map<String, FieldDescriptor> fieldsToValidate = new HashMap<>();
-        fieldsToValidate.put("Model connector name", new FieldDescriptor(mlCreateConnectorInput.getName(), true));
+
+        fieldsToValidate
+            .put("Model connector name", new FieldDescriptor(mlCreateConnectorInput.getName(), !mlCreateConnectorInput.isDryRun()));
         fieldsToValidate.put("Model connector description", new FieldDescriptor(mlCreateConnectorInput.getDescription(), false));
 
         return validateFields(fieldsToValidate);

common/src/test/java/org/opensearch/ml/common/transport/connector/MLCreateConnectorRequestTests.java

@@ -6,6 +6,7 @@
 package org.opensearch.ml.common.transport.connector;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNotSame;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertSame;
@@ -267,4 +268,47 @@ public void validateWithEmptyAndInvalidModelConnectorNameAndDescription() {
         );
     }
 
+    @Test
+    public void validateWithDryRun() {
+        // Test with dry run set to true
+        MLCreateConnectorInput dryRunInput = MLCreateConnectorInput
+            .builder()
+            .name("")  // Empty name, which would normally fail validation
+            .description("Test description")
+            .version("1")
+            .protocol("http")
+            .parameters(Map.of("input", "test"))
+            .credential(Map.of("key", "value"))
+            .actions(List.of())
+            .access(AccessMode.PUBLIC)
+            .backendRoles(Arrays.asList("role1"))
+            .addAllBackendRoles(false)
+            .dryRun(true)  // Set dry run to true
+            .build();
+
+        MLCreateConnectorRequest dryRunRequest = MLCreateConnectorRequest.builder().mlCreateConnectorInput(dryRunInput).build();
+        ActionRequestValidationException dryRunException = dryRunRequest.validate();
+        assertNull("Validation should pass when dry run is true, even with empty name", dryRunException);
+
+        // Test with dry run set to false (default behavior)
+        MLCreateConnectorInput nonDryRunInput = MLCreateConnectorInput
+            .builder()
+            .name("")  // Empty name, which should fail validation
+            .description("Test description")
+            .version("1")
+            .protocol("http")
+            .parameters(Map.of("input", "test"))
+            .credential(Map.of("key", "value"))
+            .actions(List.of())
+            .access(AccessMode.PUBLIC)
+            .backendRoles(Arrays.asList("role1"))
+            .addAllBackendRoles(false)
+            .dryRun(false)  // Set dry run to false
+            .build();
+
+        MLCreateConnectorRequest nonDryRunRequest = MLCreateConnectorRequest.builder().mlCreateConnectorInput(nonDryRunInput).build();
+        ActionRequestValidationException nonDryRunException = nonDryRunRequest.validate();
+        assertNotNull("Validation should fail when dry run is false and name is empty", nonDryRunException);
+        assertTrue(nonDryRunException.getMessage().contains("Model connector name is required"));
+    }
 }

docs/tutorials/aws/AIConnectorHelper.ipynb

@@ -37,9 +37,22 @@
     "            domain_info = response['DomainStatus']\n",
     "\n",
     "            # Extract the domain URL and ARN\n",
-    "            domain_url = domain_info['Endpoint']\n",
     "            domain_arn = domain_info['ARN']\n",
     "\n",
+    "            # Check if domain has VPC endpoints\n",
+    "            if 'Endpoints' in domain_info:\n",
+    "                # VPC domain case\n",
+    "                if 'vpc' in domain_info['Endpoints']:\n",
+    "                    domain_url = domain_info['Endpoints']['vpc']\n",
+    "                else:\n",
+    "                    domain_url = next(iter(domain_info['Endpoints'].values()))\n",
+    "            # Non-VPC domain case\n",
+    "            elif 'Endpoint' in domain_info:\n",
+    "                domain_url = domain_info['Endpoint']\n",
+    "            else:\n",
+    "                print(f\"No endpoint found for domain '{domain_name}'\")\n",
+    "                return None, None\n",
+    "\n",
     "            return f'https://{domain_url}', domain_arn\n",
     "\n",
     "        except opensearch_client.exceptions.ResourceNotFoundException:\n",

docs/tutorials/aws/DeepSeek_demo_notebook_for_RAG.ipynb

@@ -65,8 +65,22 @@
     "        try:\n",
     "            response = self.opensearch_client.describe_elasticsearch_domain(DomainName=domain_name)\n",
     "            domain_info = response['DomainStatus']\n",
-    "            domain_url = domain_info['Endpoint']\n",
     "            domain_arn = domain_info['ARN']\n",
+    "\n",
+    "            # Check if domain has VPC endpoints\n",
+    "            if 'Endpoints' in domain_info:\n",
+    "                # VPC domain case\n",
+    "                if 'vpc' in domain_info['Endpoints']:\n",
+    "                    domain_url = domain_info['Endpoints']['vpc']\n",
+    "                else:\n",
+    "                    domain_url = next(iter(domain_info['Endpoints'].values()))\n",
+    "            # Non-VPC domain case\n",
+    "            elif 'Endpoint' in domain_info:\n",
+    "                domain_url = domain_info['Endpoint']\n",
+    "            else:\n",
+    "                print(f\"No endpoint found for domain '{domain_name}'\")\n",
+    "                return None, None\n",
+    "\n",
     "            return f'https://{domain_url}', domain_arn\n",
     "        except self.opensearch_client.exceptions.ResourceNotFoundException:\n",
     "            print(f\"Domain '{domain_name}' not found.\")\n",

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLPlanExecuteAndReflectAgentRunner.java

@@ -66,6 +66,7 @@
 import org.opensearch.remote.metadata.client.SdkClient;
 import org.opensearch.transport.client.Client;
 
+import com.google.common.annotations.VisibleForTesting;
 import com.jayway.jsonpath.JsonPath;
 
 import joptsimple.internal.Strings;
@@ -154,7 +155,8 @@ public MLPlanExecuteAndReflectAgentRunner(
         this.plannerWithHistoryPromptTemplate = DEFAULT_PLANNER_WITH_HISTORY_PROMPT_TEMPLATE;
     }
 
-    private void setupPromptParameters(Map<String, String> params) {
+    @VisibleForTesting
+    void setupPromptParameters(Map<String, String> params) {
         // populated depending on whether LLM is asked to plan or re-evaluate
         // removed here, so that error is thrown in case this field is not populated
         params.remove(PROMPT_FIELD);
@@ -203,22 +205,26 @@ private void setupPromptParameters(Map<String, String> params) {
         }
     }
 
-    private void usePlannerPromptTemplate(Map<String, String> params) {
+    @VisibleForTesting
+    void usePlannerPromptTemplate(Map<String, String> params) {
         params.put(PROMPT_TEMPLATE_FIELD, this.plannerPromptTemplate);
         populatePrompt(params);
     }
 
-    private void useReflectPromptTemplate(Map<String, String> params) {
+    @VisibleForTesting
+    void useReflectPromptTemplate(Map<String, String> params) {
         params.put(PROMPT_TEMPLATE_FIELD, this.reflectPromptTemplate);
         populatePrompt(params);
     }
 
-    private void usePlannerWithHistoryPromptTemplate(Map<String, String> params) {
+    @VisibleForTesting
+    void usePlannerWithHistoryPromptTemplate(Map<String, String> params) {
         params.put(PROMPT_TEMPLATE_FIELD, this.plannerWithHistoryPromptTemplate);
         populatePrompt(params);
     }
 
-    private void populatePrompt(Map<String, String> allParams) {
+    @VisibleForTesting
+    void populatePrompt(Map<String, String> allParams) {
         String promptTemplate = allParams.get(PROMPT_TEMPLATE_FIELD);
         StringSubstitutor promptSubstitutor = new StringSubstitutor(allParams, "${parameters.", "}");
         String prompt = promptSubstitutor.replace(promptTemplate);
@@ -475,7 +481,8 @@ private void executePlanningLoop(
         client.execute(MLPredictionTaskAction.INSTANCE, request, planListener);
     }
 
-    private Map<String, String> parseLLMOutput(Map<String, String> allParams, ModelTensorOutput modelTensorOutput) {
+    @VisibleForTesting
+    Map<String, String> parseLLMOutput(Map<String, String> allParams, ModelTensorOutput modelTensorOutput) {
         Map<String, String> modelOutput = new HashMap<>();
         Map<String, ?> dataAsMap = modelTensorOutput.getMlModelOutputs().getFirst().getMlModelTensors().getFirst().getDataAsMap();
         String llmResponse;
@@ -513,7 +520,8 @@ private Map<String, String> parseLLMOutput(Map<String, String> allParams, ModelT
         return modelOutput;
     }
 
-    private String extractJsonFromMarkdown(String response) {
+    @VisibleForTesting
+    String extractJsonFromMarkdown(String response) {
         response = response.trim();
         if (response.contains("```json")) {
             response = response.substring(response.indexOf("```json") + "```json".length());
@@ -535,7 +543,8 @@ private String extractJsonFromMarkdown(String response) {
         return response;
     }
 
-    private void addToolsToPrompt(Map<String, Tool> tools, Map<String, String> allParams) {
+    @VisibleForTesting
+    void addToolsToPrompt(Map<String, Tool> tools, Map<String, String> allParams) {
         StringBuilder toolsPrompt = new StringBuilder("In this environment, you have access to the below tools: \n");
         for (Map.Entry<String, Tool> entry : tools.entrySet()) {
             String toolName = entry.getKey();
@@ -548,11 +557,13 @@ private void addToolsToPrompt(Map<String, Tool> tools, Map<String, String> allPa
         cleanUpResource(tools);
     }
 
-    private void addSteps(List<String> steps, Map<String, String> allParams, String field) {
+    @VisibleForTesting
+    void addSteps(List<String> steps, Map<String, String> allParams, String field) {
         allParams.put(field, String.join(", ", steps));
     }
 
-    private void saveAndReturnFinalResult(
+    @VisibleForTesting
+    void saveAndReturnFinalResult(
         ConversationIndexMemory memory,
         String parentInteractionId,
         String reactAgentMemoryId,
@@ -591,7 +602,8 @@ private void saveAndReturnFinalResult(
         }));
     }
 
-    private static List<ModelTensors> createModelTensors(
+    @VisibleForTesting
+    static List<ModelTensors> createModelTensors(
         String sessionId,
         String parentInteractionId,
         String reactAgentMemoryId,