From 8074492891c67a72c1584d716d738d50979622a3 Mon Sep 17 00:00:00 2001
From: xinyual <74362153+xinyual@users.noreply.github.com>
Date: Tue, 23 Jan 2024 09:38:00 +0800
Subject: [PATCH] fix name bug (#139)

* fix name bug

Signed-off-by: xinyual <xinyual@amazon.com>

* apply spotless

Signed-off-by: xinyual <xinyual@amazon.com>

* change error message

Signed-off-by: xinyual <xinyual@amazon.com>

---------

Signed-off-by: xinyual <xinyual@amazon.com>
---
 src/main/java/org/opensearch/agent/tools/PPLTool.java | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/main/java/org/opensearch/agent/tools/PPLTool.java b/src/main/java/org/opensearch/agent/tools/PPLTool.java
index 0713cb25..ce17c594 100644
--- a/src/main/java/org/opensearch/agent/tools/PPLTool.java
+++ b/src/main/java/org/opensearch/agent/tools/PPLTool.java
@@ -149,9 +149,12 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
         GetMappingsRequest getMappingsRequest = buildGetMappingRequest(indexName);
         client.admin().indices().getMappings(getMappingsRequest, ActionListener.<GetMappingsResponse>wrap(getMappingsResponse -> {
             Map<String, MappingMetadata> mappings = getMappingsResponse.getMappings();
+            if (mappings.size() == 0) {
+                throw new IllegalArgumentException("No matching mapping with index name: " + indexName);
+            }
             client.search(searchRequest, ActionListener.<SearchResponse>wrap(searchResponse -> {
                 SearchHit[] searchHits = searchResponse.getHits().getHits();
-                String tableInfo = constructTableInfo(searchHits, mappings, indexName);
+                String tableInfo = constructTableInfo(searchHits, mappings);
                 String prompt = constructPrompt(tableInfo, question, indexName);
                 RemoteInferenceInputDataSet inputDataSet = RemoteInferenceInputDataSet
                     .builder()
@@ -288,9 +291,9 @@ private GetMappingsRequest buildGetMappingRequest(String indexName) {
         return getMappingsRequest;
     }
 
-    private String constructTableInfo(SearchHit[] searchHits, Map<String, MappingMetadata> mappings, String indexName)
-        throws PrivilegedActionException {
-        MappingMetadata mappingMetadata = mappings.get(indexName);
+    private String constructTableInfo(SearchHit[] searchHits, Map<String, MappingMetadata> mappings) throws PrivilegedActionException {
+        String firstIndexName = (String) mappings.keySet().toArray()[0];
+        MappingMetadata mappingMetadata = mappings.get(firstIndexName);
         Map<String, Object> mappingSource = (Map<String, Object>) mappingMetadata.getSourceAsMap().get("properties");
         Map<String, String> fieldsToType = new HashMap<>();
         extractNamesTypes(mappingSource, fieldsToType, "");