Merge pull request #43 from VisActor/feat-customRequestFunc-skylark

Feat custom request func skylark
VisActor · Mar 19, 2024 · 19849ac · 19849ac
2 parents d64f630 + 3f7cb97
commit 19849ac
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 6 deletions.
diff --git a/packages/vmind/__tests__/browser/src/pages/ChartPreview.tsx b/packages/vmind/__tests__/browser/src/pages/ChartPreview.tsx
@@ -194,7 +194,7 @@ export function ChartPreview(props: IPropsType) {
           </div>
           {props.spec ? (
             <div>
-              <p>Total Time: {props.costTime / 1000} ms</p>
+              <p>Total Time: {props.costTime / 1000} s</p>
               <p>spec:</p>
               <TextArea value={JSON.stringify(props.spec, null, 4)} style={{ height: 300 }}></TextArea>
               {/*<pre style={{ whiteSpace: 'pre' }}>{JSON.stringify(props.spec, null, 4)}</pre>*/}

diff --git a/packages/vmind/__tests__/browser/src/pages/DataInput.tsx b/packages/vmind/__tests__/browser/src/pages/DataInput.tsx
@@ -39,7 +39,7 @@ import VMind from '../../../../src/index';
 import { Model } from '../../../../src/index';
 import { queryDataset } from '../../../../src/gpt/dataProcess';
 import { isArray } from 'lodash';
-//import { mockDataset } from './mockData'
+import { mockDataset, mockData2 } from './mockData';
 
 const TextArea = Input.TextArea;
 const Option = Select.Option;
@@ -124,7 +124,8 @@ export function DataInput(props: IPropsType) {
     const { fieldInfo, dataset } = vmind.parseCSVData(csv);
     //const { fieldInfo: fieldInfoQuery, dataset: datasetQuery } = await vmind?.dataQuery(describe, fieldInfo, dataset);
     //const { fieldInfo, dataset, usage } = await vmind.parseCSVDataWithLLM(csv, describe);
-    //const dataset = mockDataset
+
+    //const dataset = mockData2
     //const fieldInfo = vmind?.getFieldInfo(dataset)
     const startTime = new Date().getTime();
     const chartGenerationRes = await vmind.generateChart(describe, fieldInfo, dataset);

diff --git a/packages/vmind/src/core/VMind.ts b/packages/vmind/src/core/VMind.ts
@@ -108,7 +108,7 @@ class VMind {
     }
     if (this.getModelType() === ModelType.SKYLARK) {
       console.error('Please user GPT model');
-      return { fieldInfo: [], dataset };
+      return { fieldInfo: [], dataset } as any;
     }
     console.error('unsupported model in data query!');
 

diff --git a/packages/vmind/src/skylark/chart-generation/NLToChart.ts b/packages/vmind/src/skylark/chart-generation/NLToChart.ts
@@ -82,7 +82,10 @@ export const chartAdvisorSkylark = async (
     chartRecommendConstraintsStr,
     options.showThoughts ?? true
   );
-  const chartRecommendRes = await requestSkyLark(chartRecommendPrompt, userMessage, options);
+
+  const requestFunc = options.customRequestFunc?.chartAdvisor ?? requestSkyLark;
+
+  const chartRecommendRes = await requestFunc(chartRecommendPrompt, userMessage, options);
   const chartRecommendResJSON = parseSkylarkResponse(chartRecommendRes);
   if (chartRecommendResJSON.error) {
     throw Error(chartRecommendResJSON.message);
@@ -106,7 +109,7 @@ export const chartAdvisorSkylark = async (
     options.showThoughts ?? true
   );
 
-  const fieldMapRes = await requestSkyLark(fieldMapPrompt, userMessage, options);
+  const fieldMapRes = await requestFunc(fieldMapPrompt, userMessage, options);
   const fieldMapResJson = parseSkylarkResponse(fieldMapRes);
   if (fieldMapResJson.error) {
     throw Error('Network Error!');