feat: add visualization card in customized parser (#16)

* feat: add implement visualization in customized parser Signed-off-by: SuZhou-Joe <[email protected]> * feat: add csv-parser lib Signed-off-by: SuZhou-Joe <[email protected]> * feat: add csv-parser lib Signed-off-by: SuZhou-Joe <[email protected]> * feat: add test cases Signed-off-by: SuZhou-Joe <[email protected]> * feat: optimize Signed-off-by: SuZhou-Joe <[email protected]> --------- Signed-off-by: SuZhou-Joe <[email protected]>
opensearch-project · Dec 1, 2023 · 789db2b · 789db2b
1 parent b82c17a
commit 789db2b
Show file tree

Hide file tree

Showing 9 changed files with 254 additions and 3 deletions.
diff --git a/common/types/chat_saved_object_attributes.ts b/common/types/chat_saved_object_attributes.ts
@@ -12,7 +12,7 @@ export interface Interaction {
   conversation_id: string;
   interaction_id: string;
   create_time: string;
-  additional_info: Record<string, unknown>;
+  additional_info?: Record<string, unknown>;
   parent_interaction_id?: string;
 }
 

diff --git a/package.json b/package.json
@@ -19,6 +19,7 @@
   },
   "dependencies": {
     "autosize": "^6.0.1",
+    "csv-parser": "^3.0.0",
     "dompurify": "^2.4.1",
     "jsdom": "^22.1.0",
     "langchain": "^0.0.164",

diff --git a/server/parsers/basic_input_output_parser.test.ts b/server/parsers/basic_input_output_parser.test.ts
@@ -0,0 +1,32 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { BasicInputOutputParser } from './basic_input_output_parser';
+
+describe('BasicInputOutputParser', () => {
+  it('return input and output', async () => {
+    expect(
+      await BasicInputOutputParser.parserProvider({
+        input: 'input',
+        response: 'response',
+        conversation_id: '',
+        interaction_id: 'interaction_id',
+        create_time: '',
+      })
+    ).toEqual([
+      {
+        type: 'input',
+        contentType: 'text',
+        content: 'input',
+      },
+      {
+        type: 'output',
+        contentType: 'markdown',
+        content: 'response',
+        traceId: 'interaction_id',
+      },
+    ]);
+  });
+});
diff --git a/server/parsers/visualization_card_parser.test.ts b/server/parsers/visualization_card_parser.test.ts
@@ -0,0 +1,116 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { VisualizationCardParser } from './visualization_card_parser';
+
+describe('VisualizationCardParser', () => {
+  it('return visualizations when there is VisualizationTool.output', async () => {
+    expect(
+      await VisualizationCardParser.parserProvider({
+        input: 'input',
+        response: 'response',
+        conversation_id: '',
+        interaction_id: 'interaction_id',
+        create_time: '',
+        additional_info: {
+          'VisualizationTool.output': [
+            'row_number,Id,title\n' +
+              '1,id1,[Flights] Total Flights\n' +
+              '2,id2,[Flights] Controls\n' +
+              '3,id3,[Flights] Airline Carrier',
+          ],
+        },
+      })
+    ).toEqual([
+      {
+        content: 'id1',
+        contentType: 'visualization',
+        suggestedActions: [{ actionType: 'view_in_dashboards', message: 'View in Visualize' }],
+        type: 'output',
+      },
+      {
+        content: 'id2',
+        contentType: 'visualization',
+        suggestedActions: [{ actionType: 'view_in_dashboards', message: 'View in Visualize' }],
+        type: 'output',
+      },
+      {
+        content: 'id3',
+        contentType: 'visualization',
+        suggestedActions: [{ actionType: 'view_in_dashboards', message: 'View in Visualize' }],
+        type: 'output',
+      },
+    ]);
+  });
+
+  it('return visualizations when there are multiple VisualizationTool.outputs', async () => {
+    expect(
+      await VisualizationCardParser.parserProvider({
+        input: 'input',
+        response: 'response',
+        conversation_id: '',
+        interaction_id: 'interaction_id',
+        create_time: '',
+        additional_info: {
+          'VisualizationTool.output': [
+            'row_number,Id,title\n' + '1,id1,[Flights] Total Flights\n',
+            'row_number,Id,title\n' + '2,id2,[Flights] Controls\n',
+          ],
+        },
+      })
+    ).toEqual([
+      {
+        content: 'id1',
+        contentType: 'visualization',
+        suggestedActions: [{ actionType: 'view_in_dashboards', message: 'View in Visualize' }],
+        type: 'output',
+      },
+      {
+        content: 'id2',
+        contentType: 'visualization',
+        suggestedActions: [{ actionType: 'view_in_dashboards', message: 'View in Visualize' }],
+        type: 'output',
+      },
+    ]);
+  });
+
+  it('do not return visualizations when VisualizationTool.output is null', async () => {
+    expect(
+      await VisualizationCardParser.parserProvider({
+        input: 'input',
+        response: 'response',
+        conversation_id: '',
+        interaction_id: 'interaction_id',
+        create_time: '',
+        additional_info: {},
+      })
+    ).toEqual([]);
+  });
+
+  it('do not return visualizations when VisualizationTool.output is not in correct format', async () => {
+    expect(
+      await VisualizationCardParser.parserProvider({
+        input: 'input',
+        response: 'response',
+        conversation_id: '',
+        interaction_id: 'interaction_id',
+        create_time: '',
+        additional_info: {
+          'VisualizationTool.output': [
+            'row_number\n' + '1',
+            'row_number,Id,title\n' + '2,id2,[Flights] Controls\n',
+          ],
+        },
+      })
+    ).toEqual([
+      {
+        content: 'id2',
+        contentType: 'visualization',
+        suggestedActions: [{ actionType: 'view_in_dashboards', message: 'View in Visualize' }],
+        type: 'output',
+      },
+    ]);
+  });
+});
diff --git a/server/parsers/visualization_card_parser.ts b/server/parsers/visualization_card_parser.ts
@@ -0,0 +1,48 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { IMessage, Interaction } from '../../common/types/chat_saved_object_attributes';
+import { getJsonFromString } from '../utils/csv-parser-helper';
+
+const extractIdsFromCsvString = async (csv: string) => {
+  const lines = (await getJsonFromString(csv)) as Array<{ Id: string }>;
+  return lines
+    .map((line) => line.Id)
+    .filter(<T>(v: T | null | undefined): v is T => v !== null && v !== undefined);
+};
+
+export const VisualizationCardParser = {
+  id: 'core_visualization',
+  async parserProvider(interaction: Interaction) {
+    const visualizationOutputs = interaction.additional_info?.['VisualizationTool.output'] as
+      | string[]
+      | undefined;
+    if (!visualizationOutputs) {
+      return [];
+    }
+    const visualizationIds = (
+      await Promise.all(visualizationOutputs.map((output) => extractIdsFromCsvString(output)))
+    ).flatMap((id) => id);
+
+    const visOutputs: IMessage[] = visualizationIds
+      /**
+       * Empty id will be filtered
+       */
+      .filter((id) => id)
+      .map((id) => ({
+        type: 'output',
+        content: id,
+        contentType: 'visualization',
+        suggestedActions: [
+          {
+            message: 'View in Visualize',
+            actionType: 'view_in_dashboards',
+          },
+        ],
+      }));
+
+    return visOutputs;
+  },
+};
diff --git a/server/plugin.ts b/server/plugin.ts
@@ -17,13 +17,13 @@ import {
 import { OpenSearchAlertingPlugin } from './adaptors/opensearch_alerting_plugin';
 import { OpenSearchObservabilityPlugin } from './adaptors/opensearch_observability_plugin';
 import { PPLPlugin } from './adaptors/ppl_plugin';
-import { AssistantServerConfig } from './config/schema';
 import './fetch-polyfill';
 import { setupRoutes } from './routes/index';
 import { chatSavedObject } from './saved_objects/chat_saved_object';
 import { AssistantPluginSetup, AssistantPluginStart, MessageParser } from './types';
 import { chatConfigSavedObject } from './saved_objects/chat_config_saved_object';
 import { BasicInputOutputParser } from './parsers/basic_input_output_parser';
+import { VisualizationCardParser } from './parsers/visualization_card_parser';
 
 export class AssistantPlugin implements Plugin<AssistantPluginSetup, AssistantPluginStart> {
   private readonly logger: Logger;
@@ -79,6 +79,7 @@ export class AssistantPlugin implements Plugin<AssistantPluginSetup, AssistantPl
     };
 
     registerMessageParser(BasicInputOutputParser);
+    registerMessageParser(VisualizationCardParser);
 
     return {
       registerMessageParser,

diff --git a/server/utils/csv-parser-helper.test.ts b/server/utils/csv-parser-helper.test.ts
@@ -0,0 +1,21 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { getJsonFromString } from './csv-parser-helper';
+
+describe('getJsonFromString', () => {
+  it('return correct answer', async () => {
+    expect(await getJsonFromString('title,id\n1,2')).toEqual([
+      {
+        title: '1',
+        id: '2',
+      },
+    ]);
+  });
+
+  it('return empty array when string is not in correct format', async () => {
+    expect(await getJsonFromString('1,2')).toEqual([]);
+  });
+});
diff --git a/server/utils/csv-parser-helper.ts b/server/utils/csv-parser-helper.ts
@@ -0,0 +1,25 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { Readable } from 'stream';
+import csvParser from 'csv-parser';
+
+export const getJsonFromString = (
+  csvString: string,
+  options?: csvParser.Options
+): Promise<Array<Record<string, string>> | string[][]> => {
+  const results: string[][] | Array<Record<string, string>> = [];
+  return new Promise((resolve, reject) => {
+    Readable.from(csvString)
+      .pipe(csvParser(options))
+      .on('data', (data) => results.push(data))
+      .on('end', () => {
+        resolve(results);
+      })
+      .on('error', (err) => {
+        reject(err);
+      });
+  });
+};
diff --git a/yarn.lock b/yarn.lock
@@ -545,6 +545,13 @@ csstype@^3.0.2:
   resolved "https://registry.yarnpkg.com/csstype/-/csstype-3.1.2.tgz#1d4bf9d572f11c14031f0436e1c10bc1f571f50b"
   integrity sha512-I7K1Uu0MBPzaFKg4nI5Q7Vs2t+3gWWW648spaF+Rg7pI9ds18Ugn+lvg4SHczUdKlHI5LWBXyqfS8+DufyBsgQ==
 
+csv-parser@^3.0.0:
+  version "3.0.0"
+  resolved "https://registry.npmmirror.com/csv-parser/-/csv-parser-3.0.0.tgz#b88a6256d79e090a97a1b56451f9327b01d710e7"
+  integrity sha512-s6OYSXAK3IdKqYO33y09jhypG/bSDHPuyCme/IdEHfWpLf/jKcpitVFyOC6UemgGk8v7Q5u2XE0vvwmanxhGlQ==
+  dependencies:
+    minimist "^1.2.0"
+
 data-urls@^4.0.0:
   version "4.0.0"
   resolved "https://registry.yarnpkg.com/data-urls/-/data-urls-4.0.0.tgz#333a454eca6f9a5b7b0f1013ff89074c3f522dd4"
@@ -1354,7 +1361,7 @@ minimatch@^3.0.4, minimatch@^3.1.1:
   dependencies:
     brace-expansion "^1.1.7"
 
-minimist@^1.2.5, minimist@^1.2.6:
+minimist@^1.2.0, minimist@^1.2.5, minimist@^1.2.6:
   version "1.2.8"
   resolved "https://registry.yarnpkg.com/minimist/-/minimist-1.2.8.tgz#c1a464e7693302e082a075cee0c057741ac4772c"
   integrity sha512-2yyAR8qBkN3YuheJanUpWC5U3bb5osDywNB8RzDVlDwDHbocAJveqqj1u8+SVD7jkWT4yvsHCpWqqWqAxb0zCA==