add debug log

Signed-off-by: Hailong Cui <ihailong@amazon.com>

src/main/java/org/opensearch/agent/indices/IndicesHelper.java

@@ -183,7 +183,7 @@ public void addNewVectorField(String index, String modelId, ActionListener<Boole
     }
 
     private Long getDimension(String modelId) {
-        return mlClients.getEmbeddingResult(modelId, List.of("today is sunny"), mlTaskResponse -> {
+        return mlClients.getEmbeddingResult(modelId, List.of("today is sunny"), true, mlTaskResponse -> {
             ModelTensorOutput tensorOutput = (ModelTensorOutput) mlTaskResponse.getOutput();
             return tensorOutput.getMlModelOutputs().get(0).getMlModelTensors().get(0).getShape()[0];
         });

src/main/java/org/opensearch/agent/job/IndexSummaryEmbeddingJob.java

@@ -117,7 +117,7 @@ public void run() {
                     .map(sample -> (String) sample.get(INDEX_SUMMARY))
                     .collect(Collectors.toList());
 
-                List<ModelTensors> mlModelOutputs = mlClients.getEmbeddingResult(modelId, embeddingDocs, mlTaskResponse -> {
+                List<ModelTensors> mlModelOutputs = mlClients.getEmbeddingResult(modelId, embeddingDocs, true, mlTaskResponse -> {
                     ModelTensorOutput output = (ModelTensorOutput) mlTaskResponse.getOutput();
                     return output.getMlModelOutputs();
                 });
@@ -186,12 +186,12 @@ private List<Map<String, Object>> getAllIndexMappingAndSampleData() {
             indexSummaryMap.put(INDEX_PATTERNS, indexPatterns);
             indexSummaryMap.put(ALIASES, indexMetadata.getAliases().keySet());
 
-            Map<String, Object> sourceAsMap = indexMetadata.mapping().getSourceAsMap();
-            // if index don't have any mapping, ignore
-            if (sourceAsMap == null || sourceAsMap.isEmpty()) {
+            // if index have no mapping at all
+            if (indexMetadata.mapping() == null) {
                 log.debug("No mapping for index {}", indexName);
                 continue;
             }
+            Map<String, Object> sourceAsMap = indexMetadata.mapping().getSourceAsMap();
             try (XContentBuilder builder = MediaTypeRegistry.contentBuilder(MediaTypeRegistry.JSON)) {
                 builder.map(sourceAsMap);
                 String mapping = builder.toString();

src/main/java/org/opensearch/agent/job/MLClients.java

@@ -36,6 +36,7 @@
 import org.opensearch.index.query.QueryBuilders;
 import org.opensearch.index.query.TermQueryBuilder;
 import org.opensearch.ml.common.FunctionName;
+import org.opensearch.ml.common.MLTaskState;
 import org.opensearch.ml.common.agent.MLAgent;
 import org.opensearch.ml.common.agent.MLToolSpec;
 import org.opensearch.ml.common.dataset.TextDocsInputDataSet;
@@ -45,6 +46,9 @@
 import org.opensearch.ml.common.output.model.ModelResultFilter;
 import org.opensearch.ml.common.transport.MLTaskResponse;
 import org.opensearch.ml.common.transport.agent.MLSearchAgentAction;
+import org.opensearch.ml.common.transport.deploy.MLDeployModelAction;
+import org.opensearch.ml.common.transport.deploy.MLDeployModelRequest;
+import org.opensearch.ml.common.transport.deploy.MLDeployModelResponse;
 import org.opensearch.ml.common.transport.prediction.MLPredictionTaskAction;
 import org.opensearch.ml.common.transport.prediction.MLPredictionTaskRequest;
 import org.opensearch.search.SearchHit;
@@ -65,7 +69,7 @@ public MLClients(Client client, NamedXContentRegistry xContentRegistry) {
         this.xContentRegistry = xContentRegistry;
     }
 
-    public <T> T getEmbeddingResult(String modelId, List<String> texts, Function<MLTaskResponse, T> parser) {
+    public <T> T getEmbeddingResult(String modelId, List<String> texts, boolean deploy, Function<MLTaskResponse, T> parser) {
         try {
             TextDocsInputDataSet inputDataSet = TextDocsInputDataSet
                 .builder()
@@ -80,8 +84,21 @@ public <T> T getEmbeddingResult(String modelId, List<String> texts, Function<MLT
             MLTaskResponse mlTaskResponse = predictFuture.get(DEFAULT_TIMEOUT_SECOND, TimeUnit.SECONDS);
             return parser.apply(mlTaskResponse);
         } catch (Exception ex) {
+            if (deploy && ExceptionsHelper.stackTrace(ex).contains("Model not ready yet.")) {
+                log.info("Model {} has not deployed yet, try to deploy", modelId);
+                ActionFuture<MLDeployModelResponse> deployFuture = client
+                    .execute(MLDeployModelAction.INSTANCE, new MLDeployModelRequest(modelId, false));
+                try {
+                    MLDeployModelResponse mlDeployModelResponse = deployFuture.get(2, TimeUnit.MINUTES);
+                    if (mlDeployModelResponse.getStatus().equals(MLTaskState.COMPLETED.name())) {
+                        return getEmbeddingResult(modelId, texts, false, parser);
+                    }
+                } catch (Exception e) {
+                    throw ExceptionsHelper.convertToRuntime(e);
+                }
+            }
             log.error("Invoke ML embedding failed", ex);
-            throw new RuntimeException(ex);
+            throw ExceptionsHelper.convertToRuntime(ex);
         }
     }
 

src/main/java/org/opensearch/agent/tools/IndexRoutingTool.java

@@ -5,9 +5,9 @@
 
 package org.opensearch.agent.tools;
 
-import static org.opensearch.agent.job.IndexSummaryEmbeddingJob.DATA_STREAM;
-import static org.opensearch.agent.job.IndexSummaryEmbeddingJob.INDEX_NAME;
-import static org.opensearch.agent.job.IndexSummaryEmbeddingJob.INDEX_PATTERNS;
+import static org.opensearch.agent.job.IndexSummaryEmbeddingJob.DATA_STREAM_FIELD;
+import static org.opensearch.agent.job.IndexSummaryEmbeddingJob.INDEX_NAME_FIELD;
+import static org.opensearch.agent.job.IndexSummaryEmbeddingJob.INDEX_PATTERNS_FIELD;
 
 import java.io.IOException;
 import java.io.InputStream;
@@ -53,7 +53,7 @@ public class IndexRoutingTool extends VectorDBTool {
     public static final String TYPE = "IndexRoutingTool";
 
     private static final String DEFAULT_DESCRIPTION = "Use this tool to select an appropriate index for user question, "
-        + "This tool take user plain input and return list of most related indexes or `Not sure`. "
+        + "It takes 1 argument which is a string of user question and return list of most related indexes or `Not sure`. "
         + "If the tool returns `Not sure`, mark it as final answer and ask Human to provide index name";
 
     public static final int DEFAULT_K = 5;
@@ -127,6 +127,7 @@ protected Parser<SearchResponse, Object> searchResponseParser() {
 
     @Override
     public <T> void run(Map<String, String> parameters, ActionListener<T> listener) {
+        log.debug("input={}", parameters.get(INPUT_FIELD));
         // get index of knn-index
         super.run(parameters, ActionListener.wrap(res -> {
             List<Map<String, Object>> summaries = (List<Map<String, Object>>) res;
@@ -153,11 +154,13 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
             // call LLM, MLModelTool
             String question = parameters.get(INPUT_FIELD);
             String prompt = buildFinalPrompt(summaryStr, question);
+            log.debug("prompt send to inference is {}", prompt);
             // TODO use MLModelTool
             mlClients.inference(inferenceModelId, prompt, ActionListener.wrap(r -> {
                 ModelTensorOutput output = (ModelTensorOutput) r.getOutput();
                 ModelTensor modelTensor = output.getMlModelOutputs().get(0).getMlModelTensors().get(0);
                 String response = (String) modelTensor.getDataAsMap().get("response");
+                log.debug("response back from inference mode is {}", response);
                 Set<String> validIndexes = findMatchedIndex(response, summaries);
                 listener.onResponse((T) (validIndexes.isEmpty() ? "Not sure" : validIndexes.iterator().next()));
             }, exception -> { listener.onResponse((T) "Not sure"); }));
@@ -172,17 +175,18 @@ private Set<String> findMatchedIndex(String result, List<Map<String, Object>> ca
 
         Map<String, Map<String, Object>> candidateIndexMap = candidates
             .stream()
-            .collect(Collectors.toMap(m -> (String) m.get(INDEX_NAME), m -> m));
+            .collect(Collectors.toMap(m -> (String) m.get(INDEX_NAME_FIELD), m -> m));
 
         Set<String> allCandidates = candidateIndexMap.keySet();
         List<String> predictedIndexes = Arrays.stream(result.split(",")).map(String::trim).collect(Collectors.toList());
+        log.debug("all candidates are {}, predictedIndexes are {}", allCandidates, predictedIndexes);
 
         for (String predictedIndex : predictedIndexes) {
             if (allCandidates.contains(predictedIndex)) {
                 Map<String, Object> map = candidateIndexMap.get(predictedIndex);
                 // data stream back index
-                Optional<Object> dataStreamName = Optional.ofNullable(map.get(DATA_STREAM));
-                List<String> patterns = (List<String>) map.get(INDEX_PATTERNS);
+                Optional<Object> dataStreamName = Optional.ofNullable(map.get(DATA_STREAM_FIELD));
+                List<String> patterns = (List<String>) map.get(INDEX_PATTERNS_FIELD);
                 String indexPattern = getIndexPattern(patterns, predictedIndex);
                 validIndexes.add((String) dataStreamName.orElse(indexPattern));
             } else if (predictedIndex.equals("Not sure")) {

src/main/java/org/opensearch/agent/tools/utils/LLMProvider.java

@@ -17,7 +17,7 @@
 @FieldDefaults(makeFinal = true, level = AccessLevel.PRIVATE)
 public enum LLMProvider {
     OPENAI("${prompt}"),
-    ANTHROPIC("Human: ${prompt}\\nAssistant:"),
+    ANTHROPIC("\\n\\nHuman: ${prompt} \\n\\nAssistant:"),
     MISTRAL("<s>[INST] ${prompt} [/INST]"),
     NONE("${prompt}");