Add integ tests for search state api, deprovision api, and delete api

src/main/java/org/opensearch/flowframework/model/WorkflowState.java

@@ -8,12 +8,15 @@
  */
 package org.opensearch.flowframework.model;
 
+import org.opensearch.common.xcontent.LoggingDeprecationHandler;
+import org.opensearch.common.xcontent.json.JsonXContent;
 import org.opensearch.commons.authuser.User;
 import org.opensearch.core.common.ParsingException;
 import org.opensearch.core.common.io.stream.StreamInput;
 import org.opensearch.core.common.io.stream.StreamOutput;
 import org.opensearch.core.common.io.stream.Writeable;
 import org.opensearch.core.rest.RestStatus;
+import org.opensearch.core.xcontent.NamedXContentRegistry;
 import org.opensearch.core.xcontent.ToXContentObject;
 import org.opensearch.core.xcontent.XContentBuilder;
 import org.opensearch.core.xcontent.XContentParseException;
@@ -393,6 +396,22 @@ public static WorkflowState parse(XContentParser parser) throws IOException {
             .build();
     }
 
+    /**
+     * Parse a JSON workflow state
+     * @param json A string containing a JSON representation of a workflow state
+     * @return A {@link WorkflowState} represented by the JSON
+     * @throws IOException on failure to parse
+     */
+    public static WorkflowState parse(String json) throws IOException {
+        XContentParser parser = JsonXContent.jsonXContent.createParser(
+            NamedXContentRegistry.EMPTY,
+            LoggingDeprecationHandler.INSTANCE,
+            json
+        );
+        ensureExpectedToken(XContentParser.Token.START_OBJECT, parser.nextToken(), parser);
+        return parse(parser);
+    }
+
     /**
      * The workflowID associated with this workflow-state
      * @return the workflowId

src/test/java/org/opensearch/flowframework/FlowFrameworkRestTestCase.java

@@ -319,6 +319,10 @@ protected Response createWorkflow(Template template) throws Exception {
         return TestHelpers.makeRequest(client(), "POST", WORKFLOW_URI, ImmutableMap.of(), template.toJson(), null);
     }
 
+    protected Response createWorkflowWithProvision(Template template) throws Exception {
+        return TestHelpers.makeRequest(client(), "POST", WORKFLOW_URI + "?provision=true", ImmutableMap.of(), template.toJson(), null);
+    }
+
     /**
      * Helper method to invoke the Create Workflow Rest Action with dry run validation
      * @param template the template to create
@@ -364,6 +368,40 @@ protected Response provisionWorkflow(String workflowId) throws Exception {
         );
     }
 
+    /**
+     * Helper method to invoke the Deprovision Workflow Rest Action
+     * @param workflowId the workflow ID to deprovision
+     * @return a rest response
+     * @throws Exception if the request fails
+     */
+    protected Response deprovisionWorkflow(String workflowId) throws Exception {
+        return TestHelpers.makeRequest(
+            client(),
+            "POST",
+            String.format(Locale.ROOT, "%s/%s/%s", WORKFLOW_URI, workflowId, "_deprovision"),
+            ImmutableMap.of(),
+            "",
+            null
+        );
+    }
+
+    /**
+     * Helper method to invoke the Delete Workflow Rest Action
+     * @param workflowId the workflow ID to delete
+     * @return a rest response
+     * @throws Exception if the request fails
+     */
+    protected Response deleteWorkflow(String workflowId) throws Exception {
+        return TestHelpers.makeRequest(
+            client(),
+            "DELETE",
+            String.format(Locale.ROOT, "%s/%s", WORKFLOW_URI, workflowId),
+            ImmutableMap.of(),
+            "",
+            null
+        );
+    }
+
     /**
      * Helper method to invoke the Get Workflow Rest Action
      * @param workflowId the workflow ID to get the status
@@ -416,6 +454,31 @@ protected SearchResponse searchWorkflows(String query) throws Exception {
         }
     }
 
+    protected SearchResponse searchWorkflowState(String query) throws Exception {
+        Response restSearchResponse = TestHelpers.makeRequest(
+            client(),
+            "GET",
+            String.format(Locale.ROOT, "%s/state/_search", WORKFLOW_URI),
+            ImmutableMap.of(),
+            query,
+            null
+        );
+        assertEquals(RestStatus.OK, TestHelpers.restStatus(restSearchResponse));
+
+        // Parse entity content into SearchResponse
+        MediaType mediaType = MediaType.fromMediaType(restSearchResponse.getEntity().getContentType());
+        try (
+            XContentParser parser = mediaType.xContent()
+                .createParser(
+                    NamedXContentRegistry.EMPTY,
+                    DeprecationHandler.THROW_UNSUPPORTED_OPERATION,
+                    restSearchResponse.getEntity().getContent()
+                )
+        ) {
+            return SearchResponse.fromXContent(parser);
+        }
+    }
+
     /**
      * Helper method to invoke the Get Workflow Rest Action and assert the provisioning and state status
      * @param workflowId the workflow ID to get the status

src/test/java/org/opensearch/flowframework/rest/FlowFrameworkRestApiIT.java

@@ -21,10 +21,15 @@
 import org.opensearch.flowframework.model.Workflow;
 import org.opensearch.flowframework.model.WorkflowEdge;
 import org.opensearch.flowframework.model.WorkflowNode;
+import org.opensearch.flowframework.model.WorkflowState;
 
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+import java.util.stream.Collectors;
 
 import static org.opensearch.flowframework.common.CommonValue.CREDENTIAL_FIELD;
 import static org.opensearch.flowframework.common.CommonValue.PROVISION_WORKFLOW;
@@ -174,4 +179,58 @@ public void testCreateAndProvisionRemoteModelWorkflow() throws Exception {
         assertNotNull(resourcesCreated.get(0).resourceId());
     }
 
+    public void testCreateAndProvisionAgentFrameworkWorkflow() throws Exception {
+        Template template = TestHelpers.createTemplateFromFile("agent-framework.json");
+
+        // Hit Create Workflow API to create agent-framework template, with template validation check and provision parameter
+        Response response = createWorkflowWithProvision(template);
+        assertEquals(RestStatus.CREATED, TestHelpers.restStatus(response));
+        Map<String, Object> responseMap = entityAsMap(response);
+        String workflowId = (String) responseMap.get(WORKFLOW_ID);
+        // wait and ensure state is completed/done
+        assertBusy(() -> { getAndAssertWorkflowStatus(workflowId, State.COMPLETED, ProvisioningProgress.DONE); }, 30, TimeUnit.SECONDS);
+
+        // Hit Search State API with the workflow id created above
+        String query = "{\"query\":{\"ids\":{\"values\":[\"" + workflowId + "\"]}}}";
+        SearchResponse searchResponse = searchWorkflowState(query);
+        assertEquals(1, searchResponse.getHits().getTotalHits().value);
+        String searchHitSource = searchResponse.getHits().getAt(0).getSourceAsString();
+        WorkflowState searchHitWorkflowState = WorkflowState.parse(searchHitSource);
+
+        // Assert based on the agent-framework template
+        List<ResourceCreated> resourcesCreated = searchHitWorkflowState.resourcesCreated();
+        Set<String> expectedStepNames = new HashSet<>();
+        expectedStepNames.add("root_agent");
+        expectedStepNames.add("sub_agent");
+        expectedStepNames.add("openAI_connector");
+        expectedStepNames.add("gpt-3.5-model");
+        expectedStepNames.add("deployed-gpt-3.5-model");
+        Set<String> stepNames = resourcesCreated.stream().map(ResourceCreated::workflowStepId).collect(Collectors.toSet());
+
+        assertEquals(5, resourcesCreated.size());
+        assertEquals(stepNames, expectedStepNames);
+        assertNotNull(resourcesCreated.get(0).resourceId());
+
+        // Hit Deprovision API
+        Response deprovisionResponse = deprovisionWorkflow(workflowId);
+        assertEquals(RestStatus.OK, TestHelpers.restStatus(deprovisionResponse));
+        assertBusy(
+            () -> { getAndAssertWorkflowStatus(workflowId, State.NOT_STARTED, ProvisioningProgress.NOT_STARTED); },
+            30,
+            TimeUnit.SECONDS
+        );
+
+        // Hit Delete API
+        Response deleteResponse = deleteWorkflow(workflowId);
+        assertEquals(RestStatus.OK, TestHelpers.restStatus(deleteResponse));
+
+        // wait for deletion to complete
+        Thread.sleep(30000);
+
+        // Search this workflow id in global_context index to make sure it's deleted
+        SearchResponse searchResponseAfterDeletion = searchWorkflows(query);
+        assertBusy(() -> assertEquals(0, searchResponseAfterDeletion.getHits().getTotalHits().value), 30, TimeUnit.SECONDS);
+
+    }
+
 }

src/test/resources/template/agent-framework.json

@@ -0,0 +1,201 @@
+{
+  "name": "opensearch-assistant-workflow",
+  "description": "test case",
+  "use_case": "REGISTER_AGENT",
+  "version": {
+    "template": "1.0.0",
+    "compatibility": [
+      "2.12.0",
+      "3.0.0"
+    ]
+  },
+  "workflows": {
+    "provision": {
+      "nodes": [
+        {
+          "id": "openAI_connector",
+          "type": "create_connector",
+          "user_inputs": {
+            "name": "OpenAI Chat Connector",
+            "description": "The connector to public OpenAI model service for GPT 3.5",
+            "protocol": "aws_sigv4",
+            "version": "1",
+            "parameters": {
+              "region": "us-east-1",
+              "service_name": "bedrock",
+              "anthropic_version": "bedrock-2023-05-31",
+              "endpoint": "bedrock.us-east-1.amazonaws.com",
+              "auth": "Sig_V4",
+              "content_type": "application/json",
+              "max_tokens_to_sample": 8000,
+              "temperature": 0.0001
+            },
+            "credential": {
+              "access_key": "<access-key>",
+              "secret_key": "<secret-key>",
+              "session_token": "<token>"
+            },
+            "actions": [
+              {
+                "action_type": "predict",
+                "method": "POST",
+                "url": "https://bedrock-runtime.us-east-1.amazonaws.com/model/anthropic.claude-v2/invoke",
+                "request_body": "{\"prompt\":\"${parameters.prompt}\", \"max_tokens_to_sample\":${parameters.max_tokens_to_sample}, \"temperature\":${parameters.temperature},  \"anthropic_version\":\"${parameters.anthropic_version}\" }"
+              }
+            ]
+          }
+        },
+        {
+          "id": "gpt-3.5-model",
+          "type": "register_remote_model",
+          "previous_node_inputs": {
+            "openAI_connector": "connector_id"
+          },
+          "user_inputs": {
+            "name": "flow-register-remote-test-gpt-3.5-turbo",
+            "function_name": "remote",
+            "description": "test model"
+          }
+        },
+        {
+          "id": "deployed-gpt-3.5-model",
+          "type": "deploy_model",
+          "previous_node_inputs": {
+            "gpt-3.5-model": "model_id"
+          }
+        },
+        {
+          "id": "ml_model_tool",
+          "type": "create_tool",
+          "previous_node_inputs": {
+            "deployed-gpt-3.5-model": "model_id"
+          },
+          "user_inputs": {
+            "name": "MLModelTool",
+            "type": "MLModelTool",
+            "alias": "language_model_tool",
+            "description": "A general tool to answer any question. But it can't handle math problem",
+            "parameters": {
+              "prompt": "Answer the question as best you can.",
+              "response_filter": "choices[0].message.content"
+            }
+          }
+        },
+        {
+          "id": "math_tool",
+          "type": "create_tool",
+          "user_inputs": {
+            "name": "MathTool",
+            "type": "MathTool",
+            "description": "A general tool to calculate any math problem. The action input must be valid math expression, like 2+3",
+            "parameters": {
+              "max_iteration": 5
+            }
+          }
+        },
+        {
+          "id": "sub_agent",
+          "type": "register_agent",
+          "previous_node_inputs": {
+            "math_tool": "tools"
+          },
+          "user_inputs": {
+            "name": "Sub Agent",
+            "type": "cot",
+            "description": "this is a test agent",
+            "parameters": {
+              "hello": "world"
+            },
+            "llm.model_id": "ldzS04kBxRPZ5cnWrqpd",
+            "llm.parameters": {
+              "max_iteration": "5",
+              "stop_when_no_tool_found": "true"
+            },
+            "memory": {
+              "type": "conversation_index"
+            },
+            "app_type": "chatbot",
+            "created_time": 1689793598499,
+            "last_updated_time": 1689793598530
+          }
+        },
+        {
+          "id": "agent_tool",
+          "type": "create_tool",
+          "previous_node_inputs": {
+            "sub_agent": "agent_id"
+          },
+          "user_inputs": {
+            "name": "AgentTool",
+            "type": "AgentTool",
+            "description": "Root Agent Tool",
+            "parameters": {
+              "max_iteration": 5
+            }
+          }
+        },
+        {
+          "id": "root_agent",
+          "type": "register_agent",
+          "previous_node_inputs": {
+            "deployed-gpt-3.5-model": "model_id",
+            "ml_model_tool": "tools",
+            "agent_tool": "tools"
+          },
+          "user_inputs": {
+            "name": "DEMO-Test_Agent_For_CoT",
+            "type": "cot",
+            "description": "this is a test agent",
+            "parameters": {
+              "hello": "world"
+            },
+            "llm.parameters": {
+              "max_iteration": "5",
+              "stop_when_no_tool_found": "true"
+            },
+            "memory": {
+              "type": "conversation_index"
+            },
+            "app_type": "chatbot",
+            "created_time": 1689793598499,
+            "last_updated_time": 1689793598530
+          }
+        }
+      ],
+      "edges": [
+        {
+          "source": "openAI_connector",
+          "dest": "gpt-3.5-model"
+        },
+        {
+          "source": "gpt-3.5-model",
+          "dest": "deployed-gpt-3.5-model"
+        },
+        {
+          "source": "deployed-gpt-3.5-model",
+          "dest": "root_agent"
+        },
+        {
+          "source": "deployed-gpt-3.5-model",
+          "dest": "ml_model_tool"
+        },
+        {
+          "source": "ml_model_tool",
+          "dest": "root_agent"
+        },
+        {
+          "source": "math_tool",
+          "dest": "sub_agent"
+        },
+        {
+          "source": "sub_agent",
+          "dest": "agent_tool"
+        },
+        {
+          "source": "agent_tool",
+          "dest": "root_agent"
+        }
+      ]
+    }
+  }
+}