From ded29bcf239b5d51a7edbab165553ffeb8c099cc Mon Sep 17 00:00:00 2001
From: Yulong Ruan <ruanyl@amazon.com>
Date: Thu, 6 Jun 2024 21:33:20 +0800
Subject: [PATCH] handle corner case that llm returns data with extra
 `response` field

Signed-off-by: Yulong Ruan <ruanyl@amazon.com>
---
 src/plugins/vis_type_vega/public/text_to_vega.ts | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/plugins/vis_type_vega/public/text_to_vega.ts b/src/plugins/vis_type_vega/public/text_to_vega.ts
index b770488b9763..e72d140b126c 100644
--- a/src/plugins/vis_type_vega/public/text_to_vega.ts
+++ b/src/plugins/vis_type_vega/public/text_to_vega.ts
@@ -95,7 +95,12 @@ export class Text2Vega {
     const res = await this.http.post('/api/llm/text2vega', {
       body: JSON.stringify({ query }),
     });
-    const result = res.body.inference_results[0].output[0].dataAsMap;
+    let result = res.body.inference_results[0].output[0].dataAsMap;
+    // sometimes llm returns {response: <schema>} instead of <schema>
+    if (result.response) {
+      result = JSON.parse(result.response);
+    }
+    // need to escape field: geo.city -> field: geo\\.city
     escapeField(result, 'encoding');
     return result;
   }