running poc

Signed-off-by: Yulong Ruan <ruanyl@amazon.com>

src/plugins/vis_type_vega/public/components/vega_vis_editor.tsx

@@ -96,8 +96,16 @@ function VegaVisEditor({ stateParams, setValue }: VisOptionsProps<VisParams>) {
   useEffect(() => {
     const text2vega = getText2Vega();
     text2vega.getVega$().subscribe((v) => {
-      setValue('spec', JSON.stringify(v, null, 4));
+      if (!(v instanceof Error)) {
+        setValue('spec', JSON.stringify(v, null, 4));
+      }
     });
+
+    if (window['input$']) {
+      window['input$'].subscribe((v: string) => {
+        text2vega.updateInput(v);
+      });
+    }
     // text2vega.updateInput('find unique visitors and average bytes every 3 hours');
   }, []);
 

src/plugins/vis_type_vega/public/data_model/ppl_parser.ts

@@ -35,7 +35,7 @@ export class PPLQueryParser {
       const requestObject = requests.find((r) => r.dataObject.name === data.name);
 
       if (requestObject) {
-        requestObject.dataObject.values = data;
+        requestObject.dataObject.values = data.jsonData;
       }
     });
   }

src/plugins/vis_type_vega/public/text_to_vega.ts

@@ -1,26 +1,20 @@
-import { BehaviorSubject, Observable } from 'rxjs';
+import { BehaviorSubject, Observable, of } from 'rxjs';
 import {
   takeWhile,
   debounceTime,
   distinctUntilChanged,
   switchMap,
   tap,
   filter,
+  catchError,
 } from 'rxjs/operators';
 import { HttpSetup } from 'opensearch-dashboards/public';
 
 const topN = (ppl: string, n: number) => `${ppl} | head ${n}`;
 
-const t2ppl = (input: string) =>
-  new Promise<string>((resolve) =>
-    resolve(
-      'source=opensearch_dashboards_sample_data_logs | stats DISTINCT_COUNT(clientip) AS unique_visitors, AVG(bytes) AS avg_bytes by span(timestamp, 3h) AS span'
-    )
-  );
-
 const createPrompt = (input: string, ppl: string, sample: any) => {
   return `
-Your task is to generate Vega-Lite chart specifications based on the given data, the schema, the PPL query and user's instruction.
+Your task is to generate Vega-Lite chart specifications based on the given data, the data schema, the PPL query to get the data and the user's instruction.
 
 The data is represented in json format:
 ${JSON.stringify(sample.jsonData, null, 4)}
@@ -33,15 +27,15 @@ ${ppl}
 
 The user's instruction is: ${input}
 
+You seem not quite understand how to set y-scales of the different layers, the correct syntax is {"resolve": {"scale": {"y": "independent"}}} please use it when appropriate.
 Just return the chart specification json based on Vega-Lite format.
 Just reply with the json based Vega-Lite object, do not include any other content in the reply.
-Always choose a timeUnit if the data type is timestamp based on the schema.
-  `;
+`;
 };
 
 export class Text2Vega {
   input$: BehaviorSubject<string>;
-  vega$: Observable<string>;
+  vega$: Observable<string | Error>;
   http: HttpSetup;
 
   constructor(http: HttpSetup) {
@@ -54,14 +48,22 @@ export class Text2Vega {
       distinctUntilChanged(),
       // text to ppl
       switchMap(async (value) => {
-        const ppl = await t2ppl(value);
-        return {
-          input: value,
-          ppl,
-        };
+        const pplQuestion = value.split('//')[0];
+        try {
+          const ppl = await this.text2ppl(pplQuestion);
+          return {
+            input: value,
+            ppl,
+          };
+        } catch (e) {
+          return new Error('Cannot generate ppl');
+        }
       }),
       // query sample data with ppl
       switchMap(async (value) => {
+        if (value instanceof Error) {
+          return value;
+        }
         const ppl = topN(value.ppl, 5);
         const sample = await this.http.post('/api/ppl/search', {
           body: JSON.stringify({ query: ppl, format: 'jdbc' }),
@@ -70,8 +72,16 @@ export class Text2Vega {
       }),
       // call llm to generate vega
       switchMap(async (value) => {
+        if (value instanceof Error) {
+          return value;
+        }
         const prompt = createPrompt(value.input, value.ppl, value.sample);
         const result = await this.text2vega(prompt);
+        delete result.data['values'];
+        result.data.url = {
+          '%type%': 'ppl',
+          query: value.ppl,
+        };
         return result;
       })
     );
@@ -85,6 +95,24 @@ export class Text2Vega {
     return res;
   }
 
+  async text2ppl(query: string) {
+    try {
+      const pplResponse = await this.http.post('/api/llm/text2ppl', {
+        body: JSON.stringify({
+          question: query,
+          index: 'opensearch_dashboards_sample_data_logs',
+        }),
+      });
+      // eslint-disable-next-line no-console
+      console.log(pplResponse);
+      const result = JSON.parse(pplResponse.body.inference_results[0].output[0].result);
+      console.log(result);
+      return result.ppl;
+    } catch (e) {
+      console.log(e);
+    }
+  }
+
   updateInput(value: string) {
     this.input$.next(value);
   }

src/plugins/vis_type_vega/server/plugin.ts

@@ -46,7 +46,8 @@ import { setDataSourceEnabled } from './services';
 
 export const invokeText2Vega = async (
   prompt: string,
-  modelId = 'anthropic.claude-3-haiku-20240307-v1:0'
+  modelId = 'anthropic.claude-3-sonnet-20240229-v1:0'
+  // modelId = 'anthropic.claude-3-haiku-20240307-v1:0'
 ) => {
   // Create a new Bedrock Runtime client instance.
   const client = new BedrockRuntimeClient({
@@ -121,6 +122,31 @@ export class VisTypeVegaPlugin implements Plugin<VisTypeVegaPluginSetup, VisType
       })
     );
 
+    router.post(
+      {
+        path: '/api/llm/text2ppl',
+        validate: {
+          body: schema.object({
+            index: schema.string(),
+            question: schema.string(),
+          }),
+        },
+      },
+      router.handleLegacyErrors(async (context, req, res) => {
+        const result = await context.core.opensearch.client.asCurrentUser.transport.request({
+          method: 'POST',
+          path: '/_plugins/_ml/agents/uJmpgI8BGmD7E2dhTm4e/_execute',
+          body: {
+            parameters: {
+              question: req.body.question,
+              index: req.body.index,
+            },
+          },
+        });
+        return res.ok({ body: result });
+      })
+    );
+
     return {};
   }
 

src/plugins/visualize/public/application/components/visualize_top_nav.tsx

@@ -53,6 +53,7 @@ import { APP_NAME } from '../visualize_constants';
 import { getTopNavConfig } from '../utils';
 import type { IndexPattern } from '../../../../data/public';
 import chatLogo from './query_assistant_logo.svg';
+import { BehaviorSubject } from 'rxjs';
 
 interface VisualizeTopNavProps {
   currentAppState: VisualizeAppState;
@@ -71,6 +72,10 @@ interface VisualizeTopNavProps {
   onPPL?: (ppl: string) => void;
 }
 
+const input$ = new BehaviorSubject('');
+// @ts-ignore
+window['input$'] = input$;
+
 const TopNav = ({
   currentAppState,
   isChromeVisible,
@@ -224,21 +229,9 @@ const TopNav = ({
 
   const isVega = vis.type.name === 'vega';
 
-  const indexName = 'opensearch_dashboards_sample_data_ecommerce';
+  const indexName = 'opensearch_dashboards_sample_data_logs';
   const onGenerate = async () => {
-    setGenerating(true);
-    const pplResponse = await services.http.post('/api/observability/query_assist/generate_ppl', {
-      body: JSON.stringify({
-        question: value,
-        index: indexName,
-      }),
-    });
-    // eslint-disable-next-line no-console
-    console.log(pplResponse);
-    setGenerating(false);
-    if (onPPL) {
-      onPPL(pplResponse);
-    }
+    input$.next(value);
   };
 
   return isChromeVisible ? (

yarn.lock

@@ -1645,6 +1645,13 @@
   dependencies:
     regenerator-runtime "^0.14.0"
 
+"@babel/runtime@^7.23.2":
+  version "7.24.5"
+  resolved "https://registry.yarnpkg.com/@babel/runtime/-/runtime-7.24.5.tgz#230946857c053a36ccc66e1dd03b17dd0c4ed02c"
+  integrity sha512-Nms86NXrsaeU9vbBJKni6gXiEXZ4CVpYVzEjDH9Sb8vmZ3UljyA1GSOJl/6LGPO8EHLuSF9H+IxNXHPX8QHJ4g==
+  dependencies:
+    regenerator-runtime "^0.14.0"
+
 "@babel/template@^7.22.15", "@babel/template@^7.3.3":
   version "7.22.15"
   resolved "https://registry.yarnpkg.com/@babel/template/-/template-7.22.15.tgz#09576efc3830f0430f4548ef971dde1350ef2f38"
@@ -7005,6 +7012,11 @@ compare-versions@^5.0.1:
   resolved "https://registry.yarnpkg.com/compare-versions/-/compare-versions-5.0.3.tgz#a9b34fea217472650ef4a2651d905f42c28ebfd7"
   integrity sha512-4UZlZP8Z99MGEY+Ovg/uJxJuvoXuN4M6B3hKaiackiHrgzQFEe3diJi1mf1PNHbFujM7FvLrK2bpgIaImbtZ1A==
 
+complex.js@^2.1.1:
+  version "2.1.1"
+  resolved "https://registry.yarnpkg.com/complex.js/-/complex.js-2.1.1.tgz#0675dac8e464ec431fb2ab7d30f41d889fb25c31"
+  integrity sha512-8njCHOTtFFLtegk6zQo0kkVX1rngygb/KQI6z1qZxlFI3scluC+LVTCFbrkWjBv4vvLlbQ9t88IPMC6k95VTTg==
+
 component-emitter@^1.2.1, component-emitter@^1.3.0:
   version "1.3.0"
   resolved "https://registry.yarnpkg.com/component-emitter/-/component-emitter-1.3.0.tgz#16e4070fba8ae29b679f2215853ee181ab2eabc0"
@@ -7767,6 +7779,11 @@ decimal.js@^10.2.1:
   resolved "https://registry.yarnpkg.com/decimal.js/-/decimal.js-10.3.1.tgz#d8c3a444a9c6774ba60ca6ad7261c3a94fd5e783"
   integrity sha512-V0pfhfr8suzyPGOx3nmq4aHqabehUZn6Ch9kyFpV79TGDTWFmHqUqXdabR7QHqxzrYolF4+tVmJhUG4OURg5dQ==
 
+decimal.js@^10.4.3:
+  version "10.4.3"
+  resolved "https://registry.yarnpkg.com/decimal.js/-/decimal.js-10.4.3.tgz#1044092884d245d1b7f65725fa4ad4c6f781cc23"
+  integrity sha512-VBBaLc1MgL5XpzgIP7ny5Z6Nx3UrRkIViUkPUdtl9aya5amy3De1gsUUSB1g3+3sExYNjCAsAznmukyxCb1GRA==
+
 decode-uri-component@^0.2.0:
   version "0.2.2"
   resolved "https://registry.yarnpkg.com/decode-uri-component/-/decode-uri-component-0.2.2.tgz#e69dbe25d37941171dd540e024c444cd5188e1e9"
@@ -8769,6 +8786,11 @@ escalade@^3.0.2, escalade@^3.1.1:
   resolved "https://registry.yarnpkg.com/escalade/-/escalade-3.1.1.tgz#d8cfdc7000965c5a0174b4a82eaa5c0552742e40"
   integrity sha512-k0er2gUkLf8O0zKJiAhmkTnJlTvINGv7ygDNPbeIsX/TJjGJZHuh9B2UxbsaEkmlEo9MfhrSzmhIlhRlI2GXnw==
 
+escape-latex@^1.2.0:
+  version "1.2.0"
+  resolved "https://registry.yarnpkg.com/escape-latex/-/escape-latex-1.2.0.tgz#07c03818cf7dac250cce517f4fda1b001ef2bca1"
+  integrity sha512-nV5aVWW1K0wEiUIEdZ4erkGGH8mDxGyxSeqPzRNtWP7ataw+/olFObw7hujFWlVjNsaDFw5VZ5NzVSIqRgfTiw==
+
 escape-string-regexp@4.0.0, escape-string-regexp@^4.0.0:
   version "4.0.0"
   resolved "https://registry.yarnpkg.com/escape-string-regexp/-/escape-string-regexp-4.0.0.tgz#14ba83a5d373e3d311e5afca29cf5bfad965bf34"
@@ -9812,6 +9834,11 @@ fp-ts@^2.3.1:
   resolved "https://registry.yarnpkg.com/fp-ts/-/fp-ts-2.11.9.tgz#bbc204e0932954b59c98a282635754a4b624a05e"
   integrity sha512-GhYlNKkCOfdjp71ocdtyaQGoqCswEoWDJLRr+2jClnBBq2dnSOtd6QxmJdALq8UhfqCyZZ0f0lxadU4OhwY9nw==
 
+fraction.js@4.3.4:
+  version "4.3.4"
+  resolved "https://registry.yarnpkg.com/fraction.js/-/fraction.js-4.3.4.tgz#b2bac8249a610c3396106da97c5a71da75b94b1c"
+  integrity sha512-pwiTgt0Q7t+GHZA4yaLjObx4vXmmdcS0iSJ19o8d/goUGgItX9UZWKWNnLHehxviD8wU2IWRsnR8cD5+yOJP2Q==
+
 fraction.js@^4.2.0:
   version "4.2.0"
   resolved "https://registry.yarnpkg.com/fraction.js/-/fraction.js-4.2.0.tgz#448e5109a313a3527f5a3ab2119ec4cf0e0e2950"
@@ -11935,6 +11962,11 @@ jake@^10.8.5:
     filelist "^1.0.1"
     minimatch "^3.0.4"
 
+javascript-natural-sort@^0.7.1:
+  version "0.7.1"
+  resolved "https://registry.yarnpkg.com/javascript-natural-sort/-/javascript-natural-sort-0.7.1.tgz#f9e2303d4507f6d74355a73664d1440fb5a0ef59"
+  integrity sha512-nO6jcEfZWQXDhOiBtG2KvKyEptz7RVbpGP4vTD2hLBdmNQSsCiicO2Ioinv6UI4y9ukqnBpy+XZ9H6uLNgJTlw==
+
 jest-canvas-mock@^2.5.1:
   version "2.5.1"
   resolved "https://registry.yarnpkg.com/jest-canvas-mock/-/jest-canvas-mock-2.5.1.tgz#81509af658ef485e9a1bf39c64e06761517bdbcb"
@@ -13222,6 +13254,21 @@ markdown-it@^12.3.2:
     mdurl "^1.0.1"
     uc.micro "^1.0.5"
 
+mathjs@^11.8.2:
+  version "11.12.0"
+  resolved "https://registry.yarnpkg.com/mathjs/-/mathjs-11.12.0.tgz#e933e5941930d44763ddfc5bfb08b90059449b2c"
+  integrity sha512-UGhVw8rS1AyedyI55DGz9q1qZ0p98kyKPyc9vherBkoueLntPfKtPBh14x+V4cdUWK0NZV2TBwqRFlvadscSuw==
+  dependencies:
+    "@babel/runtime" "^7.23.2"
+    complex.js "^2.1.1"
+    decimal.js "^10.4.3"
+    escape-latex "^1.2.0"
+    fraction.js "4.3.4"
+    javascript-natural-sort "^0.7.1"
+    seedrandom "^3.0.5"
+    tiny-emitter "^2.1.0"
+    typed-function "^4.1.1"
+
 mathml-tag-names@^2.1.3:
   version "2.1.3"
   resolved "https://registry.yarnpkg.com/mathml-tag-names/-/mathml-tag-names-2.1.3.tgz#4ddadd67308e780cf16a47685878ee27b736a0a3"
@@ -16552,6 +16599,11 @@ secure-json-parse@^2.4.0:
   resolved "https://registry.yarnpkg.com/secure-json-parse/-/secure-json-parse-2.4.0.tgz#5aaeaaef85c7a417f76271a4f5b0cc3315ddca85"
   integrity sha512-Q5Z/97nbON5t/L/sH6mY2EacfjVGwrCcSi5D3btRO2GZ8pf1K1UN7Z9H5J57hjVU2Qzxr1xO+FmBhOvEkzCMmg==
 
+seedrandom@^3.0.5:
+  version "3.0.5"
+  resolved "https://registry.yarnpkg.com/seedrandom/-/seedrandom-3.0.5.tgz#54edc85c95222525b0c7a6f6b3543d8e0b3aa0a7"
+  integrity sha512-8OwmbklUNzwezjGInmZ+2clQmExQPvomqjL7LFqOYqtmuxRgQYqOD3mHaU+MvZn5FLUeVxVfQjwLZW/n/JFuqg==
+
 selenium-webdriver@^4.0.0-alpha.7:
   version "4.0.0-alpha.7"
   resolved "https://registry.yarnpkg.com/selenium-webdriver/-/selenium-webdriver-4.0.0-alpha.7.tgz#e3879d8457fd7ad8e4424094b7dc0540d99e6797"
@@ -17898,6 +17950,11 @@ timsort@~0.3.0:
   resolved "https://registry.yarnpkg.com/timsort/-/timsort-0.3.0.tgz#405411a8e7e6339fe64db9a234de11dc31e02bd4"
   integrity sha1-QFQRqOfmM5/mTbmiNN4R3DHgK9Q=
 
+tiny-emitter@^2.1.0:
+  version "2.1.0"
+  resolved "https://registry.yarnpkg.com/tiny-emitter/-/tiny-emitter-2.1.0.tgz#1d1a56edfc51c43e863cbb5382a72330e3555423"
+  integrity sha512-NB6Dk1A9xgQPMoGqC5CVXn123gWyte215ONT5Pp5a0yt4nlEoO1ZWeCwpncaekPHXO60i47ihFnZPiRPjRMq4Q==
+
 tiny-invariant@^1.0.2, tiny-invariant@^1.0.6:
   version "1.2.0"
   resolved "https://registry.yarnpkg.com/tiny-invariant/-/tiny-invariant-1.2.0.tgz#a1141f86b672a9148c72e978a19a73b9b94a15a9"
@@ -18284,6 +18341,11 @@ typed-array-length@^1.0.4:
     for-each "^0.3.3"
     is-typed-array "^1.1.9"
 
+typed-function@^4.1.1:
+  version "4.1.1"
+  resolved "https://registry.yarnpkg.com/typed-function/-/typed-function-4.1.1.tgz#38ce3cae31f4f513bcb263563fdad27b2afa73e8"
+  integrity sha512-Pq1DVubcvibmm8bYcMowjVnnMwPVMeh0DIdA8ad8NZY2sJgapANJmiigSUwlt+EgXxpfIv8MWrQXTIzkfYZLYQ==
+
 typedarray-to-buffer@^3.1.5:
   version "3.1.5"
   resolved "https://registry.yarnpkg.com/typedarray-to-buffer/-/typedarray-to-buffer-3.1.5.tgz#a97ee7a9ff42691b9f783ff1bc5112fe3fca9080"