From 2698f04d16c96ae3294108174f9abe617286cffe Mon Sep 17 00:00:00 2001 From: Mingshi Liu Date: Tue, 30 Jan 2024 10:27:36 -0800 Subject: [PATCH 1/4] Add IT for VectorDBTool and NeuralSparseTool Signed-off-by: Mingshi Liu --- .../integTest/BaseAgentToolsIT.java | 6 + .../org/opensearch/integTest/RAGToolIT.java | 327 ++++++++++++++++++ .../opensearch/integTest/VectorDBToolIT.java | 206 +++++++++++ ...l_with_neural_query_type_request_body.json | 24 ++ ...w_agent_of_vectordb_tool_request_body.json | 17 + ...ter_text_embedding_model_request_body.json | 5 + 6 files changed, 585 insertions(+) create mode 100644 src/test/java/org/opensearch/integTest/RAGToolIT.java create mode 100644 src/test/java/org/opensearch/integTest/VectorDBToolIT.java create mode 100644 src/test/resources/org/opensearch/agent/tools/register_flow_agent_of_ragtool_with_neural_query_type_request_body.json create mode 100644 src/test/resources/org/opensearch/agent/tools/register_flow_agent_of_vectordb_tool_request_body.json create mode 100644 src/test/resources/org/opensearch/agent/tools/register_text_embedding_model_request_body.json diff --git a/src/test/java/org/opensearch/integTest/BaseAgentToolsIT.java b/src/test/java/org/opensearch/integTest/BaseAgentToolsIT.java index f29dee0e..1bbfdc93 100644 --- a/src/test/java/org/opensearch/integTest/BaseAgentToolsIT.java +++ b/src/test/java/org/opensearch/integTest/BaseAgentToolsIT.java @@ -241,6 +241,12 @@ protected void createIndexWithConfiguration(String indexName, String indexConfig assertEquals(indexName, responseInMap.get("index").toString()); } + protected void createIngestPipelineWithConfiguration(String pipelineName, String body) throws Exception { + Response response = makeRequest(client(), "PUT", "/_ingest/pipeline/" + pipelineName, null, body, null); + Map responseInMap = parseResponseToMap(response); + assertEquals("true", responseInMap.get("acknowledged").toString()); + } + // Similar to deleteExternalIndices, but including indices with "." prefix vs. excluding them protected void deleteSystemIndices() throws IOException { final Response response = client().performRequest(new Request("GET", "/_cat/indices?format=json" + "&expand_wildcards=all")); diff --git a/src/test/java/org/opensearch/integTest/RAGToolIT.java b/src/test/java/org/opensearch/integTest/RAGToolIT.java new file mode 100644 index 00000000..dd5d749d --- /dev/null +++ b/src/test/java/org/opensearch/integTest/RAGToolIT.java @@ -0,0 +1,327 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.integTest; + +import static org.hamcrest.Matchers.allOf; +import static org.hamcrest.Matchers.containsString; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; + +import java.io.IOException; +import java.net.URISyntaxException; +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; +import java.util.Map; + +import org.junit.After; +import org.junit.Before; +import org.opensearch.client.ResponseException; + +import lombok.SneakyThrows; + +public class RAGToolIT extends BaseAgentToolsIT { + + public static String TEST_NEURAL_INDEX_NAME = "test_neural_index"; + public static String TEST_NEURAL_SPARSE_INDEX_NAME = "test_neural_sparse_index"; + private String textEmbeddingModelId; + private String sparseEncodingModelId; + private String registerAgentWithNeuralQueryRequestBody; + private String registerAgentWithNeuralSparseQueryRequestBody; + + public RAGToolIT() throws IOException, URISyntaxException {} + + @SneakyThrows + private void prepareModel() { + String requestBody = Files + .readString( + Path + .of( + this + .getClass() + .getClassLoader() + .getResource("org/opensearch/agent/tools/register_text_embedding_model_request_body.json") + .toURI() + ) + ); + textEmbeddingModelId = registerModelThenDeploy(requestBody); + + String requestBody1 = Files + .readString( + Path + .of( + this + .getClass() + .getClassLoader() + .getResource("org/opensearch/agent/tools/register_sparse_encoding_model_request_body.json") + .toURI() + ) + ); + sparseEncodingModelId = registerModelThenDeploy(requestBody1); + } + + @SneakyThrows + private void prepareIndex() { + // prepare index for neural sparse query type + createIndexWithConfiguration( + TEST_NEURAL_SPARSE_INDEX_NAME, + "{\n" + + " \"mappings\": {\n" + + " \"properties\": {\n" + + " \"text\": {\n" + + " \"type\": \"text\"\n" + + " },\n" + + " \"embedding\": {\n" + + " \"type\": \"rank_features\"\n" + + " }\n" + + " }\n" + + " }\n" + + "}" + ); + addDocToIndex( + TEST_NEURAL_SPARSE_INDEX_NAME, + "0", + List.of("text", "embedding"), + List.of("hello world", Map.of("hello", 1, "world", 2)) + ); + addDocToIndex(TEST_NEURAL_SPARSE_INDEX_NAME, "1", List.of("text", "embedding"), List.of("a b", Map.of("a", 3, "b", 4))); + + // prepare index for neural query type + String pipelineConfig = "{\n" + + " \"description\": \"text embedding pipeline\",\n" + + " \"processors\": [\n" + + " {\n" + + " \"text_embedding\": {\n" + + " \"model_id\": \"" + + textEmbeddingModelId + + "\",\n" + + " \"field_map\": {\n" + + " \"text\": \"embedding\"\n" + + " }\n" + + " }\n" + + " }\n" + + " ]\n" + + "}"; + createIngestPipelineWithConfiguration("test-embedding-model", pipelineConfig); + + String indexMapping = "{\n" + + " \"mappings\": {\n" + + " \"properties\": {\n" + + " \"text\": {\n" + + " \"type\": \"text\"\n" + + " },\n" + + " \"embedding\": {\n" + + " \"type\": \"knn_vector\",\n" + + " \"dimension\": 384,\n" + + " \"method\": {\n" + + " \"name\": \"hnsw\",\n" + + " \"space_type\": \"l2\",\n" + + " \"engine\": \"lucene\",\n" + + " \"parameters\": {\n" + + " \"ef_construction\": 128,\n" + + " \"m\": 24\n" + + " }\n" + + " }\n" + + " }\n" + + " }\n" + + " },\n" + + " \"settings\": {\n" + + " \"index\": {\n" + + " \"knn.space_type\": \"cosinesimil\",\n" + + " \"default_pipeline\": \"test-embedding-model\",\n" + + " \"knn\": \"true\"\n" + + " }\n" + + " }\n" + + "}"; + + createIndexWithConfiguration(TEST_NEURAL_INDEX_NAME, indexMapping); + + addDocToIndex(TEST_NEURAL_INDEX_NAME, "0", List.of("text"), List.of("hello world")); + + addDocToIndex(TEST_NEURAL_INDEX_NAME, "1", List.of("text"), List.of("a b")); + } + + @Before + @SneakyThrows + public void setUp() { + super.setUp(); + prepareModel(); + prepareIndex(); + String registerAgentWithNeuralQueryRequestBodyFile = Files + .readString( + Path + .of( + this + .getClass() + .getClassLoader() + .getResource( + "org/opensearch/agent/tools/register_flow_agent_of_ragtool_with_neural_query_type_request_body.json" + ) + .toURI() + ) + ); + registerAgentWithNeuralQueryRequestBody = registerAgentWithNeuralQueryRequestBodyFile + .replace("", textEmbeddingModelId) + .replace("", TEST_NEURAL_INDEX_NAME); + + registerAgentWithNeuralSparseQueryRequestBody = registerAgentWithNeuralQueryRequestBodyFile + .replace("", sparseEncodingModelId) + .replace("", TEST_NEURAL_SPARSE_INDEX_NAME) + .replace("\"query_type\": \"neural\"", "\"query_type\": \"neural_sparse\""); + } + + @After + @SneakyThrows + public void tearDown() { + super.tearDown(); + deleteExternalIndices(); + } + + public void testRAGToolWithNeuralQueryInFlowAgent() { + String agentId = createAgent(registerAgentWithNeuralQueryRequestBody); + + // neural query to test match similar text, doc1 match with higher score + String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"c\"}}"); + assertEquals( + "The agent execute response not equal with expected.", + "{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.60735726}\n" + + "{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.3785958}\n", + result + ); + + // neural query to test match exact same text case, doc0 match with higher score + String result1 = executeAgent(agentId, "{\"parameters\": {\"question\": \"hello\"}}"); + + assertEquals( + "The agent execute response not equal with expected.", + "{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.70875686}\n" + + "{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.39044854}\n", + result1 + ); + + // if blank input, call onFailure and get exception + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"\"}}")); + + org.hamcrest.MatcherAssert + .assertThat( + exception.getMessage(), + allOf(containsString("[input] is null or empty, can not process it."), containsString("illegal_argument_exception")) + ); + + } + + public void testRAGToolWithNeuralSparseQueryInFlowAgent() { + String agentId = createAgent(registerAgentWithNeuralSparseQueryRequestBody); + + // neural sparse query to test match extract same text, doc1 match with high score + String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}"); + assertEquals( + "The agent execute response not equal with expected.", + "{\"_index\":\"test_neural_sparse_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":1.2068367}\n", + result + ); + + // neural sparse query to test match extract non-existed text, no match + String result2 = executeAgent(agentId, "{\"parameters\": {\"question\": \"c\"}}"); + assertEquals("The agent execute response not equal with expected.", "Can not get any match from search result.", result2); + + // if blank input, call onFailure and get exception + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"\"}}")); + + org.hamcrest.MatcherAssert + .assertThat( + exception.getMessage(), + allOf(containsString("[input] is null or empty, can not process it."), containsString("illegal_argument_exception")) + ); + } + + public void testRAGToolWithNeuralSparseQueryInFlowAgent_withIllegalSourceField_thenGetEmptySource() { + String agentId = createAgent(registerAgentWithNeuralSparseQueryRequestBody.replace("text", "text2")); + String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}"); + assertEquals( + "The agent execute response not equal with expected.", + "{\"_index\":\"test_neural_sparse_index\",\"_source\":{},\"_id\":\"1\",\"_score\":1.2068367}\n", + result + ); + } + + public void testRAGToolWithNeuralQueryInFlowAgent_withIllegalSourceField_thenGetEmptySource() { + String agentId = createAgent(registerAgentWithNeuralQueryRequestBody.replace("text", "text2")); + String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}"); + assertEquals( + "The agent execute response not equal with expected.", + "{\"_index\":\"test_neural_index\",\"_source\":{},\"_id\":\"1\",\"_score\":0.7572355}\n" + + "{\"_index\":\"test_neural_index\",\"_source\":{},\"_id\":\"0\",\"_score\":0.38389856}\n", + result + ); + } + + public void testRAGToolWithNeuralSparseQuery_withIllegalEmbeddingField_thenThrowException() { + String agentId = createAgent(registerAgentWithNeuralSparseQueryRequestBody.replace("\"embedding\"", "\"embedding2\"")); + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); + + org.hamcrest.MatcherAssert + .assertThat( + exception.getMessage(), + allOf( + containsString("failed to create query: [neural_sparse] query only works on [rank_features] fields"), + containsString("search_phase_execution_exception") + ) + ); + } + + public void testRAGToolWithNeuralQuery_withIllegalEmbeddingField_thenThrowException() { + String agentId = createAgent(registerAgentWithNeuralQueryRequestBody.replace("\"embedding\"", "\"embedding2\"")); + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); + + org.hamcrest.MatcherAssert + .assertThat( + exception.getMessage(), + allOf( + containsString("failed to create query: Field 'embedding2' is not knn_vector type."), + containsString("query_shard_exception") + ) + ); + } + + public void testRAGToolWithNeuralSparseQueryInFlowAgent_withIllegalIndexField_thenThrowException() { + String agentId = createAgent(registerAgentWithNeuralSparseQueryRequestBody.replace(TEST_NEURAL_SPARSE_INDEX_NAME, "test_index2")); + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); + + org.hamcrest.MatcherAssert + .assertThat( + exception.getMessage(), + allOf(containsString("no such index [test_index2]"), containsString("index_not_found_exception")) + ); + } + + public void testRAGToolWithNeuralQueryInFlowAgent_withIllegalIndexField_thenThrowException() { + String agentId = createAgent(registerAgentWithNeuralQueryRequestBody.replace(TEST_NEURAL_INDEX_NAME, "test_index2")); + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); + + org.hamcrest.MatcherAssert + .assertThat( + exception.getMessage(), + allOf(containsString("no such index [test_index2]"), containsString("index_not_found_exception")) + ); + } + + public void testRAGToolWithNeuralSparseQueryInFlowAgent_withIllegalModelIdField_thenThrowException() { + String agentId = createAgent(registerAgentWithNeuralSparseQueryRequestBody.replace(sparseEncodingModelId, "test_model_id")); + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); + + org.hamcrest.MatcherAssert + .assertThat(exception.getMessage(), allOf(containsString("Failed to find model"), containsString("status_exception"))); + } + + public void testRAGToolWithNeuralQueryInFlowAgent_withIllegalModelIdField_thenThrowException() { + String agentId = createAgent(registerAgentWithNeuralQueryRequestBody.replace(textEmbeddingModelId, "test_model_id")); + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); + + org.hamcrest.MatcherAssert + .assertThat(exception.getMessage(), allOf(containsString("Failed to find model"), containsString("status_exception"))); + } +} diff --git a/src/test/java/org/opensearch/integTest/VectorDBToolIT.java b/src/test/java/org/opensearch/integTest/VectorDBToolIT.java new file mode 100644 index 00000000..5bee6b77 --- /dev/null +++ b/src/test/java/org/opensearch/integTest/VectorDBToolIT.java @@ -0,0 +1,206 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.integTest; + +import static org.hamcrest.Matchers.allOf; +import static org.hamcrest.Matchers.containsString; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertThrows; + +import java.nio.file.Files; +import java.nio.file.Path; +import java.util.List; + +import org.junit.After; +import org.junit.Before; +import org.opensearch.client.ResponseException; + +import lombok.SneakyThrows; + +public class VectorDBToolIT extends BaseAgentToolsIT { + + public static String TEST_INDEX_NAME = "test_index"; + + private String modelId; + private String registerAgentRequestBody; + + @SneakyThrows + private void prepareModel() { + String requestBody = Files + .readString( + Path + .of( + this + .getClass() + .getClassLoader() + .getResource("org/opensearch/agent/tools/register_text_embedding_model_request_body.json") + .toURI() + ) + ); + modelId = registerModelThenDeploy(requestBody); + } + + @SneakyThrows + private void prepareIndex() { + + String pipelineConfig = "{\n" + + " \"description\": \"text embedding pipeline\",\n" + + " \"processors\": [\n" + + " {\n" + + " \"text_embedding\": {\n" + + " \"model_id\": \"" + + modelId + + "\",\n" + + " \"field_map\": {\n" + + " \"text\": \"embedding\"\n" + + " }\n" + + " }\n" + + " }\n" + + " ]\n" + + "}"; + createIngestPipelineWithConfiguration("test-embedding-model", pipelineConfig); + + String indexMapping = "{\n" + + " \"mappings\": {\n" + + " \"properties\": {\n" + + " \"text\": {\n" + + " \"type\": \"text\"\n" + + " },\n" + + " \"embedding\": {\n" + + " \"type\": \"knn_vector\",\n" + + " \"dimension\": 384,\n" + + " \"method\": {\n" + + " \"name\": \"hnsw\",\n" + + " \"space_type\": \"l2\",\n" + + " \"engine\": \"lucene\",\n" + + " \"parameters\": {\n" + + " \"ef_construction\": 128,\n" + + " \"m\": 24\n" + + " }\n" + + " }\n" + + " }\n" + + " }\n" + + " },\n" + + " \"settings\": {\n" + + " \"index\": {\n" + + " \"knn.space_type\": \"cosinesimil\",\n" + + " \"default_pipeline\": \"test-embedding-model\",\n" + + " \"knn\": \"true\"\n" + + " }\n" + + " }\n" + + "}"; + + createIndexWithConfiguration(TEST_INDEX_NAME, indexMapping); + + addDocToIndex(TEST_INDEX_NAME, "0", List.of("text"), List.of("hello world")); + + addDocToIndex(TEST_INDEX_NAME, "1", List.of("text"), List.of("a b")); + } + + @Before + @SneakyThrows + public void setUp() { + super.setUp(); + prepareModel(); + prepareIndex(); + registerAgentRequestBody = Files + .readString( + Path + .of( + this + .getClass() + .getClassLoader() + .getResource("org/opensearch/agent/tools/register_flow_agent_of_vectordb_tool_request_body.json") + .toURI() + ) + ); + registerAgentRequestBody = registerAgentRequestBody.replace("", modelId); + + } + + @After + @SneakyThrows + public void tearDown() { + super.tearDown(); + deleteExternalIndices(); + } + + public void testVectorDBToolInFlowAgent() { + String agentId = createAgent(registerAgentRequestBody); + + // match similar text, doc1 match with higher score + String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"c\"}}"); + assertEquals( + "The agent execute response not equal with expected.", + "{\"_index\":\"test_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.60735726}\n" + + "{\"_index\":\"test_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.3785958}\n", + result + ); + + // match exact same text case, doc0 match with higher score + String result1 = executeAgent(agentId, "{\"parameters\": {\"question\": \"hello\"}}"); + + assertEquals( + "The agent execute response not equal with expected.", + "{\"_index\":\"test_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.70875686}\n" + + "{\"_index\":\"test_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.39044854}\n", + result1 + ); + + // if blank input, call onFailure and get exception + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"\"}}")); + + org.hamcrest.MatcherAssert + .assertThat( + exception.getMessage(), + allOf(containsString("[input] is null or empty, can not process it."), containsString("illegal_argument_exception")) + ); + } + + public void testVectorDBToolInFlowAgent_withIllegalSourceField_thenGetEmptySource() { + String agentId = createAgent(registerAgentRequestBody.replace("text", "text2")); + String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}"); + assertEquals( + "The agent execute response not equal with expected.", + "{\"_index\":\"test_index\",\"_source\":{},\"_id\":\"1\",\"_score\":0.7572355}\n" + + "{\"_index\":\"test_index\",\"_source\":{},\"_id\":\"0\",\"_score\":0.38389856}\n", + result + ); + } + + public void testVectorDBToolInFlowAgent_withIllegalEmbeddingField_thenThrowException() { + String agentId = createAgent(registerAgentRequestBody.replace("\"embedding\"", "\"embedding2\"")); + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); + + org.hamcrest.MatcherAssert + .assertThat( + exception.getMessage(), + allOf( + containsString("failed to create query: Field 'embedding2' is not knn_vector type."), + containsString("query_shard_exception") + ) + ); + } + + public void testNeuralSparseSearchToolInFlowAgent_withIllegalIndexField_thenThrowException() { + String agentId = createAgent(registerAgentRequestBody.replace("test_index", "test_index2")); + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); + + org.hamcrest.MatcherAssert + .assertThat( + exception.getMessage(), + allOf(containsString("no such index [test_index2]"), containsString("index_not_found_exception")) + ); + } + + public void testNeuralSparseSearchToolInFlowAgent_withIllegalModelIdField_thenThrowException() { + String agentId = createAgent(registerAgentRequestBody.replace(modelId, "test_model_id")); + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); + + org.hamcrest.MatcherAssert + .assertThat(exception.getMessage(), allOf(containsString("Failed to find model"), containsString("status_exception"))); + } +} diff --git a/src/test/resources/org/opensearch/agent/tools/register_flow_agent_of_ragtool_with_neural_query_type_request_body.json b/src/test/resources/org/opensearch/agent/tools/register_flow_agent_of_ragtool_with_neural_query_type_request_body.json new file mode 100644 index 00000000..86d131e8 --- /dev/null +++ b/src/test/resources/org/opensearch/agent/tools/register_flow_agent_of_ragtool_with_neural_query_type_request_body.json @@ -0,0 +1,24 @@ +{ + "name": "Test_Agent_For_RagTool", + "type": "flow", + "description": "this is a test flow agent in flow", + "tools": [ + { + "type": "RAGTool", + "description": "A description of the tool", + "parameters": { + "embedding_model_id": "", + "inference_model_id": "c8NwFo0BY4jgIz2mgLNH", + "index": "", + "embedding_field": "embedding", + "query_type": "neural", + "enable_Content_Generation":"false", + "source_field": [ + "text" + ], + "input": "${parameters.question}", + "prompt": "\n\nHuman:You are a professional data analysist. You will always answer question based on the given context first. If the answer is not directly shown in the context, you will analyze the data and find the answer. If you don't know the answer, just say don't know. \n\n Context:\n${parameters.output_field}\n\nHuman:${parameters.question}\n\nAssistant:" + } + } + ] +} \ No newline at end of file diff --git a/src/test/resources/org/opensearch/agent/tools/register_flow_agent_of_vectordb_tool_request_body.json b/src/test/resources/org/opensearch/agent/tools/register_flow_agent_of_vectordb_tool_request_body.json new file mode 100644 index 00000000..3b13e443 --- /dev/null +++ b/src/test/resources/org/opensearch/agent/tools/register_flow_agent_of_vectordb_tool_request_body.json @@ -0,0 +1,17 @@ +{ + "name": "Test_VectorDB_Agent", + "type": "flow", + "tools": [ + { + "type": "VectorDBTool", + "parameters": { + "description":"user this tool to search data from the test index", + "model_id": "", + "index": "test_index", + "embedding_field": "embedding", + "source_field": ["text"], + "input": "${parameters.question}" + } + } + ] +} \ No newline at end of file diff --git a/src/test/resources/org/opensearch/agent/tools/register_text_embedding_model_request_body.json b/src/test/resources/org/opensearch/agent/tools/register_text_embedding_model_request_body.json new file mode 100644 index 00000000..f40412f5 --- /dev/null +++ b/src/test/resources/org/opensearch/agent/tools/register_text_embedding_model_request_body.json @@ -0,0 +1,5 @@ +{ + "name": "huggingface/sentence-transformers/all-MiniLM-L12-v2", + "version": "1.0.1", + "model_format": "TORCH_SCRIPT" +} \ No newline at end of file From d82afe0fac55303da973f697a26012918d4cb00c Mon Sep 17 00:00:00 2001 From: Mingshi Liu Date: Fri, 16 Feb 2024 17:41:08 -0800 Subject: [PATCH 2/4] changed to smaller models and refine IT tests Signed-off-by: Mingshi Liu --- .../org/opensearch/agent/tools/RAGTool.java | 2 +- .../org/opensearch/integTest/RAGToolIT.java | 231 +++++++++++++++--- .../opensearch/integTest/VectorDBToolIT.java | 26 +- ...l_with_neural_query_type_request_body.json | 5 +- ...ter_text_embedding_model_request_body.json | 15 +- 5 files changed, 228 insertions(+), 51 deletions(-) diff --git a/src/main/java/org/opensearch/agent/tools/RAGTool.java b/src/main/java/org/opensearch/agent/tools/RAGTool.java index 6c341b05..e8159839 100644 --- a/src/main/java/org/opensearch/agent/tools/RAGTool.java +++ b/src/main/java/org/opensearch/agent/tools/RAGTool.java @@ -56,7 +56,7 @@ public class RAGTool implements Tool { public static final String EMBEDDING_FIELD = "embedding_field"; public static final String OUTPUT_FIELD = "output_field"; public static final String QUERY_TYPE = "query_type"; - public static final String CONTENT_GENERATION_FIELD = "enable_Content_Generation"; + public static final String CONTENT_GENERATION_FIELD = "enable_content_generation"; public static final String K_FIELD = "k"; private final AbstractRetrieverTool queryTool; private String name = TYPE; diff --git a/src/test/java/org/opensearch/integTest/RAGToolIT.java b/src/test/java/org/opensearch/integTest/RAGToolIT.java index dd5d749d..7bf09aa7 100644 --- a/src/test/java/org/opensearch/integTest/RAGToolIT.java +++ b/src/test/java/org/opensearch/integTest/RAGToolIT.java @@ -19,18 +19,48 @@ import org.junit.After; import org.junit.Before; +import org.opensearch.agent.tools.RAGTool; import org.opensearch.client.ResponseException; import lombok.SneakyThrows; -public class RAGToolIT extends BaseAgentToolsIT { +public class RAGToolIT extends ToolIntegrationTest { public static String TEST_NEURAL_INDEX_NAME = "test_neural_index"; public static String TEST_NEURAL_SPARSE_INDEX_NAME = "test_neural_sparse_index"; private String textEmbeddingModelId; private String sparseEncodingModelId; + private String largeLanguageModelId; private String registerAgentWithNeuralQueryRequestBody; private String registerAgentWithNeuralSparseQueryRequestBody; + private String registerAgentWithNeuralQueryAndLLMRequestBody; + private String mockLLMResponseWithSource = "{\n" + + " \"inference_results\": [\n" + + " {\n" + + " \"output\": [\n" + + " {\n" + + " \"name\": \"response\",\n" + + " \"result\": \"\"\" Based on the context given:\n" + + " a, b, c are alphabets.\"\"\"\n" + + " }\n" + + " ]\n" + + " }\n" + + " ]\n" + + "}"; + private String mockLLMResponseWithoutSource = "{\n" + + " \"inference_results\": [\n" + + " {\n" + + " \"output\": [\n" + + " {\n" + + " \"name\": \"response\",\n" + + " \"result\": \"\"\" Based on the context given:\n" + + " I do not see any information about a, b, c\". So I would have to say I don't know the answer to your question based on this context..\"\"\"\n" + + " }\n" + + " ]\n" + + " }\n" + + " ]\n" + + "}"; + private String registerAgentWithNeuralSparseQueryAndLLMRequestBody; public RAGToolIT() throws IOException, URISyntaxException {} @@ -61,6 +91,7 @@ private void prepareModel() { ) ); sparseEncodingModelId = registerModelThenDeploy(requestBody1); + largeLanguageModelId = this.modelId; } @SneakyThrows @@ -115,7 +146,7 @@ private void prepareIndex() { + " },\n" + " \"embedding\": {\n" + " \"type\": \"knn_vector\",\n" - + " \"dimension\": 384,\n" + + " \"dimension\": 768,\n" + " \"method\": {\n" + " \"name\": \"hnsw\",\n" + " \"space_type\": \"l2\",\n" @@ -171,6 +202,16 @@ public void setUp() { .replace("", sparseEncodingModelId) .replace("", TEST_NEURAL_SPARSE_INDEX_NAME) .replace("\"query_type\": \"neural\"", "\"query_type\": \"neural_sparse\""); + + registerAgentWithNeuralQueryAndLLMRequestBody = registerAgentWithNeuralQueryRequestBodyFile + .replace("", textEmbeddingModelId + "\" ,\n \"inference_model_id\": \"" + largeLanguageModelId) + .replace("", TEST_NEURAL_INDEX_NAME) + .replace("false", "true"); + registerAgentWithNeuralSparseQueryAndLLMRequestBody = registerAgentWithNeuralQueryRequestBodyFile + .replace("", sparseEncodingModelId + "\" ,\n \"inference_model_id\": \"" + largeLanguageModelId) + .replace("", TEST_NEURAL_SPARSE_INDEX_NAME) + .replace("\"query_type\": \"neural\"", "\"query_type\": \"neural_sparse\"") + .replace("false", "true"); } @After @@ -178,17 +219,19 @@ public void setUp() { public void tearDown() { super.tearDown(); deleteExternalIndices(); + deleteModel(textEmbeddingModelId); + deleteModel(sparseEncodingModelId); } - public void testRAGToolWithNeuralQueryInFlowAgent() { + public void testRAGToolWithNeuralQuery() { String agentId = createAgent(registerAgentWithNeuralQueryRequestBody); // neural query to test match similar text, doc1 match with higher score String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"c\"}}"); assertEquals( "The agent execute response not equal with expected.", - "{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.60735726}\n" - + "{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.3785958}\n", + "{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.7046764}\n" + + "{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.2649903}\n", result ); @@ -197,8 +240,8 @@ public void testRAGToolWithNeuralQueryInFlowAgent() { assertEquals( "The agent execute response not equal with expected.", - "{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.70875686}\n" - + "{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.39044854}\n", + "{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.56714886}\n" + + "{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.24236833}\n", result1 ); @@ -208,12 +251,30 @@ public void testRAGToolWithNeuralQueryInFlowAgent() { org.hamcrest.MatcherAssert .assertThat( exception.getMessage(), - allOf(containsString("[input] is null or empty, can not process it."), containsString("illegal_argument_exception")) + allOf(containsString("[input] is null or empty, can not process it."), containsString("IllegalArgumentException")) + ); + + } + + public void testRAGToolWithNeuralQueryAndLLM() { + String agentId = createAgent(registerAgentWithNeuralQueryAndLLMRequestBody); + + // neural query to test match similar text, doc1 match with higher score + String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"use RAGTool to answer c\"}}"); + assertEquals(mockLLMResponseWithSource, result); + + // if blank input, call onFailure and get exception + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"\"}}")); + + org.hamcrest.MatcherAssert + .assertThat( + exception.getMessage(), + allOf(containsString("[input] is null or empty, can not process it."), containsString("IllegalArgumentException")) ); } - public void testRAGToolWithNeuralSparseQueryInFlowAgent() { + public void testRAGToolWithNeuralSparseQuery() { String agentId = createAgent(registerAgentWithNeuralSparseQueryRequestBody); // neural sparse query to test match extract same text, doc1 match with high score @@ -234,11 +295,28 @@ public void testRAGToolWithNeuralSparseQueryInFlowAgent() { org.hamcrest.MatcherAssert .assertThat( exception.getMessage(), - allOf(containsString("[input] is null or empty, can not process it."), containsString("illegal_argument_exception")) + allOf(containsString("[input] is null or empty, can not process it."), containsString("IllegalArgumentException")) + ); + } + + public void testRAGToolWithNeuralSparseQueryAndLLM() { + String agentId = createAgent(registerAgentWithNeuralSparseQueryAndLLMRequestBody); + + // neural sparse query to test match extract same text, doc1 match with high score + String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"use RAGTool to answer a\"}}"); + assertEquals(mockLLMResponseWithSource, result); + + // if blank input, call onFailure and get exception + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"\"}}")); + + org.hamcrest.MatcherAssert + .assertThat( + exception.getMessage(), + allOf(containsString("[input] is null or empty, can not process it."), containsString("IllegalArgumentException")) ); } - public void testRAGToolWithNeuralSparseQueryInFlowAgent_withIllegalSourceField_thenGetEmptySource() { + public void testRAGToolWithNeuralSparseQuery_withIllegalSourceField_thenGetEmptySource() { String agentId = createAgent(registerAgentWithNeuralSparseQueryRequestBody.replace("text", "text2")); String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}"); assertEquals( @@ -248,17 +326,29 @@ public void testRAGToolWithNeuralSparseQueryInFlowAgent_withIllegalSourceField_t ); } - public void testRAGToolWithNeuralQueryInFlowAgent_withIllegalSourceField_thenGetEmptySource() { + public void testRAGToolWithNeuralSparseQueryAndLLM_withIllegalSourceField_thenGetEmptySource() { + String agentId = createAgent(registerAgentWithNeuralSparseQueryAndLLMRequestBody.replace("text", "text2")); + String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}"); + assertEquals(mockLLMResponseWithoutSource, result); + } + + public void testRAGToolWithNeuralQuery_withIllegalSourceField_thenGetEmptySource() { String agentId = createAgent(registerAgentWithNeuralQueryRequestBody.replace("text", "text2")); String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}"); assertEquals( "The agent execute response not equal with expected.", - "{\"_index\":\"test_neural_index\",\"_source\":{},\"_id\":\"1\",\"_score\":0.7572355}\n" - + "{\"_index\":\"test_neural_index\",\"_source\":{},\"_id\":\"0\",\"_score\":0.38389856}\n", + "{\"_index\":\"test_neural_index\",\"_source\":{},\"_id\":\"0\",\"_score\":0.70493275}\n" + + "{\"_index\":\"test_neural_index\",\"_source\":{},\"_id\":\"1\",\"_score\":0.2650575}\n", result ); } + public void testRAGToolWithNeuralQueryAndLLM_withIllegalSourceField_thenGetEmptySource() { + String agentId = createAgent(registerAgentWithNeuralQueryAndLLMRequestBody.replace("text", "text2")); + String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}"); + assertEquals(mockLLMResponseWithoutSource, result); + } + public void testRAGToolWithNeuralSparseQuery_withIllegalEmbeddingField_thenThrowException() { String agentId = createAgent(registerAgentWithNeuralSparseQueryRequestBody.replace("\"embedding\"", "\"embedding2\"")); Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); @@ -266,10 +356,18 @@ public void testRAGToolWithNeuralSparseQuery_withIllegalEmbeddingField_thenThrow org.hamcrest.MatcherAssert .assertThat( exception.getMessage(), - allOf( - containsString("failed to create query: [neural_sparse] query only works on [rank_features] fields"), - containsString("search_phase_execution_exception") - ) + allOf(containsString("all shards failed"), containsString("SearchPhaseExecutionException")) + ); + } + + public void testRAGToolWithNeuralSparseQueryAndLLM_withIllegalEmbeddingField_thenThrowException() { + String agentId = createAgent(registerAgentWithNeuralSparseQueryAndLLMRequestBody.replace("\"embedding\"", "\"embedding2\"")); + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); + + org.hamcrest.MatcherAssert + .assertThat( + exception.getMessage(), + allOf(containsString("all shards failed"), containsString("SearchPhaseExecutionException")) ); } @@ -280,48 +378,121 @@ public void testRAGToolWithNeuralQuery_withIllegalEmbeddingField_thenThrowExcept org.hamcrest.MatcherAssert .assertThat( exception.getMessage(), - allOf( - containsString("failed to create query: Field 'embedding2' is not knn_vector type."), - containsString("query_shard_exception") - ) + allOf(containsString("all shards failed"), containsString("SearchPhaseExecutionException")) + ); + } + + public void testRAGToolWithNeuralQueryAndLLM_withIllegalEmbeddingField_thenThrowException() { + String agentId = createAgent(registerAgentWithNeuralQueryAndLLMRequestBody.replace("\"embedding\"", "\"embedding2\"")); + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); + + org.hamcrest.MatcherAssert + .assertThat( + exception.getMessage(), + allOf(containsString("all shards failed"), containsString("SearchPhaseExecutionException")) ); } - public void testRAGToolWithNeuralSparseQueryInFlowAgent_withIllegalIndexField_thenThrowException() { + public void testRAGToolWithNeuralSparseQuery_withIllegalIndexField_thenThrowException() { String agentId = createAgent(registerAgentWithNeuralSparseQueryRequestBody.replace(TEST_NEURAL_SPARSE_INDEX_NAME, "test_index2")); Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); org.hamcrest.MatcherAssert .assertThat( exception.getMessage(), - allOf(containsString("no such index [test_index2]"), containsString("index_not_found_exception")) + allOf(containsString("no such index [test_index2]"), containsString("IndexNotFoundException")) + ); + } + + public void testRAGToolWithNeuralSparseQueryAndLLM_withIllegalIndexField_thenThrowException() { + String agentId = createAgent( + registerAgentWithNeuralSparseQueryAndLLMRequestBody.replace(TEST_NEURAL_SPARSE_INDEX_NAME, "test_index2") + ); + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); + + org.hamcrest.MatcherAssert + .assertThat( + exception.getMessage(), + allOf(containsString("no such index [test_index2]"), containsString("IndexNotFoundException")) ); } - public void testRAGToolWithNeuralQueryInFlowAgent_withIllegalIndexField_thenThrowException() { + public void testRAGToolWithNeuralQuery_withIllegalIndexField_thenThrowException() { String agentId = createAgent(registerAgentWithNeuralQueryRequestBody.replace(TEST_NEURAL_INDEX_NAME, "test_index2")); Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); org.hamcrest.MatcherAssert .assertThat( exception.getMessage(), - allOf(containsString("no such index [test_index2]"), containsString("index_not_found_exception")) + allOf(containsString("no such index [test_index2]"), containsString("IndexNotFoundException")) + ); + } + + public void testRAGToolWithNeuralQueryAndLLM_withIllegalIndexField_thenThrowException() { + String agentId = createAgent(registerAgentWithNeuralQueryAndLLMRequestBody.replace(TEST_NEURAL_INDEX_NAME, "test_index2")); + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); + + org.hamcrest.MatcherAssert + .assertThat( + exception.getMessage(), + allOf(containsString("no such index [test_index2]"), containsString("IndexNotFoundException")) ); } - public void testRAGToolWithNeuralSparseQueryInFlowAgent_withIllegalModelIdField_thenThrowException() { + public void testRAGToolWithNeuralSparseQuery_withIllegalModelIdField_thenThrowException() { String agentId = createAgent(registerAgentWithNeuralSparseQueryRequestBody.replace(sparseEncodingModelId, "test_model_id")); Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); org.hamcrest.MatcherAssert - .assertThat(exception.getMessage(), allOf(containsString("Failed to find model"), containsString("status_exception"))); + .assertThat(exception.getMessage(), allOf(containsString("Failed to find model"), containsString("OpenSearchStatusException"))); + } + + public void testRAGToolWithNeuralSparseQueryAndLLM_withIllegalModelIdField_thenThrowException() { + String agentId = createAgent(registerAgentWithNeuralSparseQueryAndLLMRequestBody.replace(sparseEncodingModelId, "test_model_id")); + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); + + org.hamcrest.MatcherAssert + .assertThat(exception.getMessage(), allOf(containsString("Failed to find model"), containsString("OpenSearchStatusException"))); } - public void testRAGToolWithNeuralQueryInFlowAgent_withIllegalModelIdField_thenThrowException() { + public void testRAGToolWithNeuralQuery_withIllegalModelIdField_thenThrowException() { String agentId = createAgent(registerAgentWithNeuralQueryRequestBody.replace(textEmbeddingModelId, "test_model_id")); Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); org.hamcrest.MatcherAssert - .assertThat(exception.getMessage(), allOf(containsString("Failed to find model"), containsString("status_exception"))); + .assertThat(exception.getMessage(), allOf(containsString("Failed to find model"), containsString("OpenSearchStatusException"))); + } + + public void testRAGToolWithNeuralQueryAndLLM_withIllegalModelIdField_thenThrowException() { + String agentId = createAgent(registerAgentWithNeuralQueryAndLLMRequestBody.replace(textEmbeddingModelId, "test_model_id")); + Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); + + org.hamcrest.MatcherAssert + .assertThat(exception.getMessage(), allOf(containsString("Failed to find model"), containsString("OpenSearchStatusException"))); + } + + @Override + List promptHandlers() { + PromptHandler RAGToolHandler = new PromptHandler() { + @Override + String response(String prompt) { + if (prompt.contains("RAGTool")) { + return mockLLMResponseWithSource; + } else { + return mockLLMResponseWithoutSource; + } + } + + @Override + boolean apply(String prompt) { + return true; + } + }; + return List.of(RAGToolHandler); + } + + @Override + String toolType() { + return RAGTool.TYPE; } } diff --git a/src/test/java/org/opensearch/integTest/VectorDBToolIT.java b/src/test/java/org/opensearch/integTest/VectorDBToolIT.java index 5bee6b77..6a1db2fc 100644 --- a/src/test/java/org/opensearch/integTest/VectorDBToolIT.java +++ b/src/test/java/org/opensearch/integTest/VectorDBToolIT.java @@ -71,7 +71,7 @@ private void prepareIndex() { + " },\n" + " \"embedding\": {\n" + " \"type\": \"knn_vector\",\n" - + " \"dimension\": 384,\n" + + " \"dimension\": 768,\n" + " \"method\": {\n" + " \"name\": \"hnsw\",\n" + " \"space_type\": \"l2\",\n" @@ -126,6 +126,7 @@ public void setUp() { public void tearDown() { super.tearDown(); deleteExternalIndices(); + deleteModel(modelId); } public void testVectorDBToolInFlowAgent() { @@ -135,8 +136,8 @@ public void testVectorDBToolInFlowAgent() { String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"c\"}}"); assertEquals( "The agent execute response not equal with expected.", - "{\"_index\":\"test_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.60735726}\n" - + "{\"_index\":\"test_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.3785958}\n", + "{\"_index\":\"test_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.7046764}\n" + + "{\"_index\":\"test_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.2649903}\n", result ); @@ -145,8 +146,8 @@ public void testVectorDBToolInFlowAgent() { assertEquals( "The agent execute response not equal with expected.", - "{\"_index\":\"test_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.70875686}\n" - + "{\"_index\":\"test_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.39044854}\n", + "{\"_index\":\"test_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.56714886}\n" + + "{\"_index\":\"test_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.24236833}\n", result1 ); @@ -156,7 +157,7 @@ public void testVectorDBToolInFlowAgent() { org.hamcrest.MatcherAssert .assertThat( exception.getMessage(), - allOf(containsString("[input] is null or empty, can not process it."), containsString("illegal_argument_exception")) + allOf(containsString("[input] is null or empty, can not process it."), containsString("IllegalArgumentException")) ); } @@ -165,8 +166,8 @@ public void testVectorDBToolInFlowAgent_withIllegalSourceField_thenGetEmptySourc String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}"); assertEquals( "The agent execute response not equal with expected.", - "{\"_index\":\"test_index\",\"_source\":{},\"_id\":\"1\",\"_score\":0.7572355}\n" - + "{\"_index\":\"test_index\",\"_source\":{},\"_id\":\"0\",\"_score\":0.38389856}\n", + "{\"_index\":\"test_index\",\"_source\":{},\"_id\":\"0\",\"_score\":0.70493275}\n" + + "{\"_index\":\"test_index\",\"_source\":{},\"_id\":\"1\",\"_score\":0.2650575}\n", result ); } @@ -178,10 +179,7 @@ public void testVectorDBToolInFlowAgent_withIllegalEmbeddingField_thenThrowExcep org.hamcrest.MatcherAssert .assertThat( exception.getMessage(), - allOf( - containsString("failed to create query: Field 'embedding2' is not knn_vector type."), - containsString("query_shard_exception") - ) + allOf(containsString("all shards failed"), containsString("SearchPhaseExecutionException")) ); } @@ -192,7 +190,7 @@ public void testNeuralSparseSearchToolInFlowAgent_withIllegalIndexField_thenThro org.hamcrest.MatcherAssert .assertThat( exception.getMessage(), - allOf(containsString("no such index [test_index2]"), containsString("index_not_found_exception")) + allOf(containsString("no such index [test_index2]"), containsString("IndexNotFoundException")) ); } @@ -201,6 +199,6 @@ public void testNeuralSparseSearchToolInFlowAgent_withIllegalModelIdField_thenTh Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); org.hamcrest.MatcherAssert - .assertThat(exception.getMessage(), allOf(containsString("Failed to find model"), containsString("status_exception"))); + .assertThat(exception.getMessage(), allOf(containsString("Failed to find model"), containsString("OpenSearchStatusException"))); } } diff --git a/src/test/resources/org/opensearch/agent/tools/register_flow_agent_of_ragtool_with_neural_query_type_request_body.json b/src/test/resources/org/opensearch/agent/tools/register_flow_agent_of_ragtool_with_neural_query_type_request_body.json index 86d131e8..7f8a4af2 100644 --- a/src/test/resources/org/opensearch/agent/tools/register_flow_agent_of_ragtool_with_neural_query_type_request_body.json +++ b/src/test/resources/org/opensearch/agent/tools/register_flow_agent_of_ragtool_with_neural_query_type_request_body.json @@ -8,16 +8,15 @@ "description": "A description of the tool", "parameters": { "embedding_model_id": "", - "inference_model_id": "c8NwFo0BY4jgIz2mgLNH", "index": "", "embedding_field": "embedding", "query_type": "neural", - "enable_Content_Generation":"false", + "enable_content_generation":"false", "source_field": [ "text" ], "input": "${parameters.question}", - "prompt": "\n\nHuman:You are a professional data analysist. You will always answer question based on the given context first. If the answer is not directly shown in the context, you will analyze the data and find the answer. If you don't know the answer, just say don't know. \n\n Context:\n${parameters.output_field}\n\nHuman:${parameters.question}\n\nAssistant:" + "prompt": "\n\nHuman:You are a professional data analyst. You will always answer question based on the given context first. If the answer is not directly shown in the context, you will analyze the data and find the answer. If you don't know the answer, just say don't know. \n\n Context:\n${parameters.output_field}\n\nHuman:${parameters.question}\n\nAssistant:" } } ] diff --git a/src/test/resources/org/opensearch/agent/tools/register_text_embedding_model_request_body.json b/src/test/resources/org/opensearch/agent/tools/register_text_embedding_model_request_body.json index f40412f5..0173665a 100644 --- a/src/test/resources/org/opensearch/agent/tools/register_text_embedding_model_request_body.json +++ b/src/test/resources/org/opensearch/agent/tools/register_text_embedding_model_request_body.json @@ -1,5 +1,14 @@ { - "name": "huggingface/sentence-transformers/all-MiniLM-L12-v2", - "version": "1.0.1", - "model_format": "TORCH_SCRIPT" + "name": "traced_small_model", + "version": "1.0.0", + "model_format": "TORCH_SCRIPT", + "model_task_type": "text_embedding", + "model_content_hash_value": "e13b74006290a9d0f58c1376f9629d4ebc05a0f9385f40db837452b167ae9021", + "model_config": { + "model_type": "bert", + "embedding_dimension": 768, + "framework_type": "sentence_transformers", + "all_config": "{\"architectures\":[\"BertModel\"],\"max_position_embeddings\":512,\"model_type\":\"bert\",\"num_attention_heads\":12,\"num_hidden_layers\":6}" + }, + "url": "https://github.com/opensearch-project/ml-commons/blob/2.x/ml-algorithms/src/test/resources/org/opensearch/ml/engine/algorithms/text_embedding/traced_small_model.zip?raw=true" } \ No newline at end of file From 9f8049a1dc9a6ef78e5f2f057b7e6b8d7ce26bd2 Mon Sep 17 00:00:00 2001 From: Mingshi Liu Date: Sun, 3 Mar 2024 20:20:38 -0800 Subject: [PATCH 3/4] add match full prompt Signed-off-by: Mingshi Liu --- .../org/opensearch/integTest/RAGToolIT.java | 23 +++++++++++++++++-- .../opensearch/integTest/VectorDBToolIT.java | 4 ++-- 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/src/test/java/org/opensearch/integTest/RAGToolIT.java b/src/test/java/org/opensearch/integTest/RAGToolIT.java index 7bf09aa7..989f60eb 100644 --- a/src/test/java/org/opensearch/integTest/RAGToolIT.java +++ b/src/test/java/org/opensearch/integTest/RAGToolIT.java @@ -260,7 +260,7 @@ public void testRAGToolWithNeuralQueryAndLLM() { String agentId = createAgent(registerAgentWithNeuralQueryAndLLMRequestBody); // neural query to test match similar text, doc1 match with higher score - String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"use RAGTool to answer c\"}}"); + String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"use RAGTool to answer a\"}}"); assertEquals(mockLLMResponseWithSource, result); // if blank input, call onFailure and get exception @@ -476,7 +476,26 @@ List promptHandlers() { PromptHandler RAGToolHandler = new PromptHandler() { @Override String response(String prompt) { - if (prompt.contains("RAGTool")) { + String expectPromptForNeuralSparseQuery = "\n" + + "\nHuman:You are a professional data analyst. You will always answer question based on the given context first. If the answer is not directly shown in the context, you will analyze the data and find the answer. If you don't know the answer, just say don't know. \n" + + "\n" + + " Context:\n" + + "\"_id: 1\\n_source: {\\\"text\\\":\\\"a b\\\"}\\n\"\n" + + "\n" + + "Human:use RAGTool to answer a\n" + + "\n" + + "Assistant:"; + String expectPromptForNeuralQuery = "\n" + + "\n" + + "Human:You are a professional data analyst. You will always answer question based on the given context first. If the answer is not directly shown in the context, you will analyze the data and find the answer. If you don't know the answer, just say don't know. \n" + + "\n" + + " Context:\n" + + "\"_id: 1\\n_source: {\\\"text\\\":\\\"a b\\\"}\\n_id: 0\\n_source: {\\\"text\\\":\\\"hello world\\\"}\\n\"\n" + + "\n" + + "Human:use RAGTool to answer a\n" + + "\n" + + "Assistant:"; + if (prompt.equals(expectPromptForNeuralSparseQuery) || prompt.equals(expectPromptForNeuralQuery)) { return mockLLMResponseWithSource; } else { return mockLLMResponseWithoutSource; diff --git a/src/test/java/org/opensearch/integTest/VectorDBToolIT.java b/src/test/java/org/opensearch/integTest/VectorDBToolIT.java index 6a1db2fc..6f1314d1 100644 --- a/src/test/java/org/opensearch/integTest/VectorDBToolIT.java +++ b/src/test/java/org/opensearch/integTest/VectorDBToolIT.java @@ -183,7 +183,7 @@ public void testVectorDBToolInFlowAgent_withIllegalEmbeddingField_thenThrowExcep ); } - public void testNeuralSparseSearchToolInFlowAgent_withIllegalIndexField_thenThrowException() { + public void testVectorDBToolInFlowAgent_withIllegalIndexField_thenThrowException() { String agentId = createAgent(registerAgentRequestBody.replace("test_index", "test_index2")); Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); @@ -194,7 +194,7 @@ public void testNeuralSparseSearchToolInFlowAgent_withIllegalIndexField_thenThro ); } - public void testNeuralSparseSearchToolInFlowAgent_withIllegalModelIdField_thenThrowException() { + public void testVectorDBToolInFlowAgent_withIllegalModelIdField_thenThrowException() { String agentId = createAgent(registerAgentRequestBody.replace(modelId, "test_model_id")); Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}")); From e0483d97f1073421ff1c3b8edc8dcb30bc62c9ef Mon Sep 17 00:00:00 2001 From: Mingshi Liu Date: Tue, 5 Mar 2024 15:30:10 -0800 Subject: [PATCH 4/4] fix digits rounding differences in CI Signed-off-by: Mingshi Liu --- .../org/opensearch/integTest/RAGToolIT.java | 30 +++++++++--------- .../opensearch/integTest/VectorDBToolIT.java | 31 ++++++++----------- 2 files changed, 27 insertions(+), 34 deletions(-) diff --git a/src/test/java/org/opensearch/integTest/RAGToolIT.java b/src/test/java/org/opensearch/integTest/RAGToolIT.java index 989f60eb..a444e7dd 100644 --- a/src/test/java/org/opensearch/integTest/RAGToolIT.java +++ b/src/test/java/org/opensearch/integTest/RAGToolIT.java @@ -228,22 +228,21 @@ public void testRAGToolWithNeuralQuery() { // neural query to test match similar text, doc1 match with higher score String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"c\"}}"); - assertEquals( - "The agent execute response not equal with expected.", - "{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.7046764}\n" - + "{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.2649903}\n", - result + + // To allow digits variation from model output, using string contains to match + assertTrue( + result.contains("{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.70467") ); + assertTrue(result.contains("{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.26499")); // neural query to test match exact same text case, doc0 match with higher score String result1 = executeAgent(agentId, "{\"parameters\": {\"question\": \"hello\"}}"); - assertEquals( - "The agent execute response not equal with expected.", - "{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.56714886}\n" - + "{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.24236833}\n", - result1 + // To allow digits variation from model output, using string contains to match + assertTrue( + result1.contains("{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.5671488") ); + assertTrue(result1.contains("{\"_index\":\"test_neural_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.2423683")); // if blank input, call onFailure and get exception Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"\"}}")); @@ -335,12 +334,11 @@ public void testRAGToolWithNeuralSparseQueryAndLLM_withIllegalSourceField_thenGe public void testRAGToolWithNeuralQuery_withIllegalSourceField_thenGetEmptySource() { String agentId = createAgent(registerAgentWithNeuralQueryRequestBody.replace("text", "text2")); String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}"); - assertEquals( - "The agent execute response not equal with expected.", - "{\"_index\":\"test_neural_index\",\"_source\":{},\"_id\":\"0\",\"_score\":0.70493275}\n" - + "{\"_index\":\"test_neural_index\",\"_source\":{},\"_id\":\"1\",\"_score\":0.2650575}\n", - result - ); + + // To allow digits variation from model output, using string contains to match + assertTrue(result.contains("{\"_index\":\"test_neural_index\",\"_source\":{},\"_id\":\"0\",\"_score\":0.70493")); + assertTrue(result.contains("{\"_index\":\"test_neural_index\",\"_source\":{},\"_id\":\"1\",\"_score\":0.26505")); + } public void testRAGToolWithNeuralQueryAndLLM_withIllegalSourceField_thenGetEmptySource() { diff --git a/src/test/java/org/opensearch/integTest/VectorDBToolIT.java b/src/test/java/org/opensearch/integTest/VectorDBToolIT.java index 6f1314d1..668564dc 100644 --- a/src/test/java/org/opensearch/integTest/VectorDBToolIT.java +++ b/src/test/java/org/opensearch/integTest/VectorDBToolIT.java @@ -7,7 +7,6 @@ import static org.hamcrest.Matchers.allOf; import static org.hamcrest.Matchers.containsString; -import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertThrows; import java.nio.file.Files; @@ -134,22 +133,19 @@ public void testVectorDBToolInFlowAgent() { // match similar text, doc1 match with higher score String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"c\"}}"); - assertEquals( - "The agent execute response not equal with expected.", - "{\"_index\":\"test_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.7046764}\n" - + "{\"_index\":\"test_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.2649903}\n", - result - ); + + // To allow digits variation from model output, using string contains to match + assertTrue(result.contains("{\"_index\":\"test_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.70467")); + assertTrue(result.contains("{\"_index\":\"test_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.26499")); // match exact same text case, doc0 match with higher score String result1 = executeAgent(agentId, "{\"parameters\": {\"question\": \"hello\"}}"); - assertEquals( - "The agent execute response not equal with expected.", - "{\"_index\":\"test_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.56714886}\n" - + "{\"_index\":\"test_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.24236833}\n", - result1 + // To allow digits variation from model output, using string contains to match + assertTrue( + result1.contains("{\"_index\":\"test_index\",\"_source\":{\"text\":\"hello world\"},\"_id\":\"0\",\"_score\":0.5671488") ); + assertTrue(result1.contains("{\"_index\":\"test_index\",\"_source\":{\"text\":\"a b\"},\"_id\":\"1\",\"_score\":0.2423683")); // if blank input, call onFailure and get exception Exception exception = assertThrows(ResponseException.class, () -> executeAgent(agentId, "{\"parameters\": {\"question\": \"\"}}")); @@ -164,12 +160,11 @@ public void testVectorDBToolInFlowAgent() { public void testVectorDBToolInFlowAgent_withIllegalSourceField_thenGetEmptySource() { String agentId = createAgent(registerAgentRequestBody.replace("text", "text2")); String result = executeAgent(agentId, "{\"parameters\": {\"question\": \"a\"}}"); - assertEquals( - "The agent execute response not equal with expected.", - "{\"_index\":\"test_index\",\"_source\":{},\"_id\":\"0\",\"_score\":0.70493275}\n" - + "{\"_index\":\"test_index\",\"_source\":{},\"_id\":\"1\",\"_score\":0.2650575}\n", - result - ); + + // To allow digits variation from model output, using string contains to match + assertTrue(result.contains("{\"_index\":\"test_index\",\"_source\":{},\"_id\":\"0\",\"_score\":0.70493")); + assertTrue(result.contains("{\"_index\":\"test_index\",\"_source\":{},\"_id\":\"1\",\"_score\":0.26505")); + } public void testVectorDBToolInFlowAgent_withIllegalEmbeddingField_thenThrowException() {