add chat extraction and change return key

Signed-off-by: xinyual <[email protected]>
xinyual · Dec 11, 2023 · 9ff84a7 · 9ff84a7
1 parent 2b13cb3
commit 9ff84a7
Showing 1 changed file with 16 additions and 1 deletion.
diff --git a/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/PPLTool.java b/ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/PPLTool.java
@@ -94,6 +94,7 @@ public PPLTool(Client client, String modelId, String contextPrompt) {
 
     @Override
     public <T> void run(Map<String, String> parameters, ActionListener<T> listener) {
+        parameters = extractFromChatParameters(parameters);
         String indexName = parameters.get("index");
         String question = parameters.get("question");
         SearchRequest searchRequest = buildSearchRequest(indexName);
@@ -111,7 +112,7 @@ 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();
-                    String ppl = dataAsMap.get("output");
+                    String ppl = dataAsMap.get("response");
                     JSONObject jsonContent = new JSONObject(ImmutableMap.of("query", ppl));
                     PPLQueryRequest pplQueryRequest = new PPLQueryRequest(ppl, jsonContent, null, "jdbc");
                     TransportPPLQueryRequest transportPPLQueryRequest = new TransportPPLQueryRequest(pplQueryRequest);
@@ -313,4 +314,18 @@ private static TransportPPLQueryResponse fromActionResponse(ActionResponse actio
         }
 
     }
+
+    private Map<String, String> extractFromChatParameters(Map<String, String> parameters)
+    {
+        if (parameters.containsKey("input"))
+        {
+            try {
+                Map<String, String> chatParameters = gson.fromJson(parameters.get("input"), Map.class);
+                parameters.putAll(chatParameters);
+            } finally {
+                return parameters;
+            }
+        }
+        return parameters;
+    }
 }