Deal endpoint failure problem (#326)

* deal endpoint failure problem Signed-off-by: xinyual <[email protected]> * deal endpoint failure problem Signed-off-by: xinyual <[email protected]> * add ut Signed-off-by: xinyual <[email protected]> * apply spotless Signed-off-by: xinyual <[email protected]> * change exception type Signed-off-by: xinyual <[email protected]> --------- Signed-off-by: xinyual <[email protected]>
opensearch-project · Jun 13, 2024 · f14418b · f14418b
1 parent 5362a9a
commit f14418b
Show file tree

Hide file tree

Showing 2 changed files with 40 additions and 0 deletions.
diff --git a/src/main/java/org/opensearch/agent/tools/PPLTool.java b/src/main/java/org/opensearch/agent/tools/PPLTool.java
@@ -222,6 +222,10 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
                     ModelTensors modelTensors = modelTensorOutput.getMlModelOutputs().get(0);
                     ModelTensor modelTensor = modelTensors.getMlModelTensors().get(0);
                     Map<String, String> dataAsMap = (Map<String, String>) modelTensor.getDataAsMap();
+                    if (dataAsMap.get("response") == null) {
+                        listener.onFailure(new IllegalStateException("Remote endpoint fails to inference."));
+                        return;
+                    }
                     String ppl = parseOutput(dataAsMap.get("response"), indexName);
                     if (!this.execute) {
                         Map<String, String> ret = ImmutableMap.of("ppl", ppl);

diff --git a/src/test/java/org/opensearch/agent/tools/PPLToolTests.java b/src/test/java/org/opensearch/agent/tools/PPLToolTests.java
@@ -265,6 +265,42 @@ public void testTool_withPPLTag() {
 
     }
 
+    @Test
+    public void testTool_withWrongEndpointInference() {
+        PPLTool tool = PPLTool.Factory.getInstance().create(ImmutableMap.of("model_id", "modelId", "prompt", "contextPrompt"));
+        assertEquals(PPLTool.TYPE, tool.getName());
+
+        pplReturns = Collections.singletonMap("code", "424");
+        modelTensor = new ModelTensor("tensor", new Number[0], new long[0], MLResultDataType.STRING, null, null, pplReturns);
+        initMLTensors();
+
+        Exception exception = assertThrows(
+            IllegalStateException.class,
+            () -> tool.run(ImmutableMap.of("index", "demo", "question", "demo"), ActionListener.<String>wrap(ppl -> {
+                assertEquals(pplResult, "ppl result");
+            }, e -> { throw new IllegalStateException(e.getMessage()); }))
+        );
+        assertEquals("Remote endpoint fails to inference.", exception.getMessage());
+    }
+
+    @Test
+    public void testTool_withWrongEndpointInferenceWithNullResponse() {
+        PPLTool tool = PPLTool.Factory.getInstance().create(ImmutableMap.of("model_id", "modelId", "prompt", "contextPrompt"));
+        assertEquals(PPLTool.TYPE, tool.getName());
+
+        pplReturns = Collections.singletonMap("response", null);
+        modelTensor = new ModelTensor("tensor", new Number[0], new long[0], MLResultDataType.STRING, null, null, pplReturns);
+        initMLTensors();
+
+        Exception exception = assertThrows(
+            IllegalStateException.class,
+            () -> tool.run(ImmutableMap.of("index", "demo", "question", "demo"), ActionListener.<String>wrap(ppl -> {
+                assertEquals(pplResult, "ppl result");
+            }, e -> { throw new IllegalStateException(e.getMessage()); }))
+        );
+        assertEquals("Remote endpoint fails to inference.", exception.getMessage());
+    }
+
     @Test
     public void testTool_withDescribeStartPPL() {
         PPLTool tool = PPLTool.Factory.getInstance().create(ImmutableMap.of("model_id", "modelId", "prompt", "contextPrompt"));