feat: refactor advise (#836)

* feat(ava): add new chart advise prompt

* fix(ava): fix type

* fix(ava): fix render bug

* fix(ava): fix ci problem

---------

Co-authored-by: xueren.dt <xueren.dt@antgroup.com>

Author: leondt1

packages/ava/src/advisor/advise-chart-pipeline/plugins/advisePlugin.ts

@@ -1,24 +1,19 @@
-import { logError, logInDev, isOpenAi, isTbox, requestLLM, safeJsonParse } from '@ava/utils';
+import { logError, isOpenAi, isTbox, requestLLM, safeJsonParse } from '@ava/utils';
 import {
   AdviseChartParams,
   AdviseChartPluginInput,
   AdvisorPlugin,
-  ChartConfig,
-  AdviseChart,
   IAdviseChartPipeline,
   DataShard,
   PlainLikeDataType,
 } from '@ava/types';
-import {
-  generateAllChartConfigs,
-  optimizeChartConfig,
-  sortChartConfigs,
-  transformChartEncode,
-} from '@ava/advisor/chartAdvise';
-import { getPlainChartAdvisePrompt } from '@ava/advisor/chartAdvise/prompt';
 import { AdviseChartPluginEnum } from '@ava/constants/pipeline';
 import { DATA_SHAPE } from '@ava/extract/constants';
 
+import { getChartAdvisePrompt, getSpecGeneratePrompt } from '../../chartAdvise/prompt';
+import { CHART_ID_MAP } from '../../../ckb';
+import { Spec } from '../../../bind';
+
 export class AdvisePlugin implements AdvisorPlugin<AdviseChartParams> {
   name = AdviseChartPluginEnum.AdvisePlugin;
 
@@ -41,86 +36,52 @@ export class AdvisePlugin implements AdvisorPlugin<AdviseChartParams> {
   };
 
   advisePlain = async (dataShards: DataShard[], input: AdviseChartPluginInput) => {
-    let allChartConfigs: Array<ChartConfig[]> = [];
-    let _llmCompleted = false;
-    let _llmCostTime = '';
     const { context, dataStore } = input;
-    const { excludes, includes, disableModel, forceType, llm, uiConfig = {} } = context;
-    const paramList = dataShards.map((shard) => {
-      const { data, metas, purpose } = shard;
-      // generate all valid chart configs using field data
-      const curConfigs = generateAllChartConfigs(metas, excludes, includes);
-      allChartConfigs.push(curConfigs);
-      return {
-        userInput: purpose?.purposeDesc ?? '',
-        chartConfig: curConfigs,
-        metas,
-        data: data as PlainLikeDataType,
-      };
-    });
-    if (forceType && dataShards.length === 1) {
-      // user specified chart type
-      const finalRes: AdviseChart[] = allChartConfigs[0]
-        .filter((item) => item.type === forceType)
-        .map((item) => ({
-          type: item.type,
-          encode: transformChartEncode(item.encode),
-        }));
-
-      const result = [
-        {
-          adviseCharts: finalRes,
-          metas: dataShards[0].metas,
-          data: dataShards[0].data as PlainLikeDataType,
-        },
-      ];
-      dataStore.advise = result;
+    const { llm } = context;
+    if (!isOpenAi(llm) && !isTbox(llm)) {
+      logError('LLM config is missing or invalid');
+      dataStore.advise = [];
     } else {
-      if (!disableModel) {
-        try {
-          logInDev.debug('All possible chart configs', JSON.stringify(allChartConfigs));
-          // Use LLM to score all chart configs based on user purpose/data/metas
-          const prompt = getPlainChartAdvisePrompt(paramList);
-          const startTime = performance.now();
-          let LLMRes = '';
-          if (!isOpenAi(llm) && !isTbox(llm)) {
-            logError('LLM config is missing or invalid');
-          } else {
-            LLMRes = await requestLLM({ config: llm, prompt });
-          }
-          if (LLMRes) {
-            logInDev.debug('chart configs after LLM scoring', LLMRes);
-            const LLMResArr = safeJsonParse(LLMRes, []);
-            if (LLMResArr.length === allChartConfigs.length) {
-              allChartConfigs = allChartConfigs.map((configs, index) => sortChartConfigs(configs, LLMResArr[index]));
-            }
-            const endTime = performance.now();
-            _llmCostTime = ((endTime - startTime) / 1000).toFixed(2);
-            _llmCompleted = true;
-          } else {
-            logError('LLM scoring failed');
-          }
-        } catch (error) {
-          logError('LLM scoring failed', error);
-        }
-      }
-      // Optimize chart configuration based on rules
-      const finalRes = allChartConfigs.map((configs, index) => {
-        const { data, metas } = dataShards[index];
-        const adviseCharts = optimizeChartConfig({
-          chartConfigs: configs,
-          metas,
-          data: data as PlainLikeDataType,
-          uiConfig,
+      try {
+        const params = dataShards.map((shard) => {
+          const { data, metas, purpose } = shard;
+          return {
+            metas,
+            data: data as PlainLikeDataType,
+            purpose: purpose?.purposeDesc ?? '',
+          };
+        });
+        const adviseInputs = getChartAdvisePrompt(params);
+        const recommendationStr = await requestLLM({ config: llm, prompt: adviseInputs });
+        const recommendation = safeJsonParse(recommendationStr, []);
+        const bestCharts = recommendation.map((item) => CHART_ID_MAP[item[0]]) as string[];
+        const specGenerateInputs = dataShards.map((shard, index) => {
+          const { data } = shard;
+          return {
+            data: data as PlainLikeDataType,
+            chartId: bestCharts[index],
+          };
         });
-        return {
-          adviseCharts,
-          metas,
-          data: data as PlainLikeDataType,
-        };
-      });
-      logInDev.debug('chart configs after optimization', JSON.stringify(finalRes));
-      dataStore.advise = finalRes;
+        const chartSpecGeneratePrompt = getSpecGeneratePrompt(specGenerateInputs);
+        const chartSpecsStr = await requestLLM({ config: llm, prompt: chartSpecGeneratePrompt });
+        const chartSpecs = safeJsonParse(chartSpecsStr, []) as Spec[];
+        const res = chartSpecs.map((item, index) => ({
+          charts: [
+            {
+              spec: {
+                ...item,
+                type: bestCharts[index],
+              },
+            },
+          ],
+          data: dataShards[index].data as PlainLikeDataType,
+          metas: dataShards[index].metas,
+        }));
+        dataStore.advise = res;
+      } catch (error) {
+        logError('LLM request failed');
+        dataStore.advise = [];
+      }
     }
   };
 

packages/ava/src/advisor/advisor.ts

@@ -1,3 +1,4 @@
+/* eslint-disable no-dupe-class-members */
 import {
   AdviseChartParams,
   AdviseStageOutput,
@@ -10,14 +11,14 @@ import { extractData } from '@ava/extract';
 
 import { logError } from '../utils';
 import { getRenderer, type Spec } from '../bind';
+
 import { AdviseChartPipeline } from './advise-chart-pipeline/pipeline';
 
 /**
  * The 1st level advisor class.
  * Used to extract data, recommend charts, and render charts. All functionalities are AI-driven.
  */
 export class Advisor {
-
   /**
    * Configuration for the advisor. Includes LLM settings, chart inclusion/exclusion lists.
    */
@@ -36,13 +37,13 @@ export class Advisor {
 
   /**
    * Extract data shards from raw data or user query.
-   * 
+   *
    * Case 1 - User provides a text query, extract relevant data shards:
-   * 
+   *
       advisor.extract('What is the average age of people who work as engineers?');
-   * 
+   *
    * Case 2 - User provides a query with raw data:
-   * 
+   *
       advisor.extract(`帮我可视化以下数据：
         城市 类别 渠道 销售额 价格
         杭州 体育 A  100  80
@@ -51,9 +52,9 @@ export class Advisor {
         广州 体育 B  120  70
         深圳 体育 C  180  95
       `);
-   * 
+   *
    * Case 3 - User provides raw data, extract data shards directly:
-   * 
+   *
       advisor.extract({ type: 'A', value: 2 });
    */
   async extract(params: AdviseChartParams) {
@@ -66,40 +67,41 @@ export class Advisor {
   /**
    * Advise charts based on the data shards, which are extracted from the `advisor.extract` API.
    * This is the core function of the Advisor class, which leverages LLMs to recommend suitable chart types and encodings based on the provided data and user purpose.
-   * 
+   *
    * const advises = advisor.advise(dataShards);
    */
   advise(params: AdviseChartParams): Promise<AdviseStageOutput>;
+
   advise(params: AdviseTextParams): Promise<AdviseText>;
+
   async advise(params: AdviseChartParams | AdviseTextParams): Promise<AdviseStageOutput | AdviseText> {
     const result = await this.adviseChartPipeline.execute(params as AdviseChartParams);
     return result;
   }
 
   /**
    * Render the chart recommendation spec into a chart dom. We can use different renderers by customizing the renderer.
-   * 
+   *
    * Case 1 - Using default renderer (AVA built-in):
-   * 
+   *
       advisor.render({
         container: '#chart',
         spec: chartSpec,
       });
-   * 
+   *
    * Case 2 - Using custom renderer:
-   * 
+   *
       const customRenderer: Renderer = (container, spec) => { ... };
       bindRenderer(customRenderer);
       advisor.render({
         container: '#chart',
         spec: chartSpec,
       });
-   * 
+   *
    */
-  render(params: { container: string; spec: Spec }) {
+  render(container: string, spec: Spec) {
     const renderer = getRenderer();
     if (renderer) {
-      const { container, spec } = params;
       return renderer(container, spec);
     }
     logError('Chart render not configured, please bind a renderer first, GPT-Vis is recommended.');