From 5ffaf3c0f131f7fbf47583f38e5948ccdf9ed249 Mon Sep 17 00:00:00 2001
From: Yulong Ruan <ruanyl@amazon.com>
Date: Fri, 17 May 2024 17:30:16 +0800
Subject: [PATCH] deploy t2vega model with ml-commons

Signed-off-by: Yulong Ruan <ruanyl@amazon.com>
---
 .../vis_type_vega/public/text_to_vega.ts      | 37 ++++++++++---------
 src/plugins/vis_type_vega/server/plugin.ts    | 14 ++++++-
 .../components/visualize_top_nav.tsx          |  6 +++
 3 files changed, 38 insertions(+), 19 deletions(-)

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 ad654f4c76a6..d6d11810c7a1 100644
--- a/src/plugins/vis_type_vega/public/text_to_vega.ts
+++ b/src/plugins/vis_type_vega/public/text_to_vega.ts
@@ -1,13 +1,5 @@
-import { BehaviorSubject, Observable, of } from 'rxjs';
-import {
-  takeWhile,
-  debounceTime,
-  distinctUntilChanged,
-  switchMap,
-  tap,
-  filter,
-  catchError,
-} from 'rxjs/operators';
+import { BehaviorSubject, Observable } from 'rxjs';
+import { debounceTime, distinctUntilChanged, switchMap, tap, filter } from 'rxjs/operators';
 import { HttpSetup } from 'opensearch-dashboards/public';
 
 const topN = (ppl: string, n: number) => `${ppl} | head ${n}`;
@@ -33,6 +25,10 @@ Just reply with the json based Vega-Lite object, do not include any other conten
 `;
 };
 
+const llmRunning$ = new BehaviorSubject(false);
+// @ts-ignore
+window['llmRunning$'] = llmRunning$;
+
 export class Text2Vega {
   input$: BehaviorSubject<string>;
   vega$: Observable<string | Error>;
@@ -42,10 +38,10 @@ export class Text2Vega {
     this.http = http;
     this.input$ = new BehaviorSubject('');
     this.vega$ = this.input$.pipe(
-      tap((v) => console.log(v)),
       filter((v) => v.length > 0),
       debounceTime(200),
       distinctUntilChanged(),
+      tap((v) => llmRunning$.next(true)),
       // text to ppl
       switchMap(async (value) => {
         const pplQuestion = value.split('//')[0];
@@ -83,16 +79,23 @@ export class Text2Vega {
           query: value.ppl,
         };
         return result;
-      })
+      }),
+      tap(() => llmRunning$.next(false))
     );
   }
 
   async text2vega(query: string) {
-    const res = await this.http.post('/api/llm/text2vega', {
-      body: JSON.stringify({ query }),
-    });
-    console.log('llm res: ', res);
-    return res;
+    try {
+      const res = await this.http.post('/api/llm/text2vega', {
+        body: JSON.stringify({ query }),
+      });
+      console.log('llm res: ', res);
+      // return res;
+      const result = res.body.inference_results[0].output[0].dataAsMap;
+      return result;
+    } catch (e) {
+      console.log(e);
+    }
   }
 
   async text2ppl(query: string) {
diff --git a/src/plugins/vis_type_vega/server/plugin.ts b/src/plugins/vis_type_vega/server/plugin.ts
index c204d336ce37..c27757c644ba 100644
--- a/src/plugins/vis_type_vega/server/plugin.ts
+++ b/src/plugins/vis_type_vega/server/plugin.ts
@@ -44,7 +44,7 @@ import {
 } from './vega_visualization_client_wrapper';
 import { setDataSourceEnabled } from './services';
 
-export const invokeText2Vega = async (
+const invokeText2Vega = async (
   prompt: string,
   modelId = 'anthropic.claude-3-sonnet-20240229-v1:0'
   // modelId = 'anthropic.claude-3-haiku-20240307-v1:0'
@@ -117,7 +117,17 @@ export class VisTypeVegaPlugin implements Plugin<VisTypeVegaPluginSetup, VisType
         },
       },
       router.handleLegacyErrors(async (context, req, res) => {
-        const result = await invokeText2Vega(req.body.query);
+        // const result = await invokeText2Vega(req.body.query);
+        // return res.ok({ body: result });
+        const result = await context.core.opensearch.client.asCurrentUser.transport.request({
+          method: 'POST',
+          path: '/_plugins/_ml/models/_yV0hY8B8ef_5QXJp6Xd/_predict',
+          body: {
+            parameters: {
+              prompt: req.body.query,
+            },
+          },
+        });
         return res.ok({ body: result });
       })
     );
diff --git a/src/plugins/visualize/public/application/components/visualize_top_nav.tsx b/src/plugins/visualize/public/application/components/visualize_top_nav.tsx
index a18ff41f0b43..0c0884431863 100644
--- a/src/plugins/visualize/public/application/components/visualize_top_nav.tsx
+++ b/src/plugins/visualize/public/application/components/visualize_top_nav.tsx
@@ -222,6 +222,12 @@ const TopNav = ({
   const [isPopoverOpen, setIsPopoverOpen] = useState(false);
   const inputRef = useRef<HTMLInputElement>(null);
 
+  useEffect(() => {
+    window['llmRunning$'].subscribe((running) => {
+      setGenerating(!!running);
+    });
+  }, []);
+
   const HARDCODED_SUGGESTIONS: string[] = [
     `what's the revenue for past week and group by day?`,
     'how many orders per day for past week?',