|
|
import path from "path"; |
|
|
import fs from "fs"; |
|
|
import puppeteer from "puppeteer"; |
|
|
import VMind, { ChartType, DataTable } from "@visactor/vmind"; |
|
|
import { isString } from "@visactor/vutils"; |
|
|
|
|
|
enum AlgorithmType { |
|
|
OverallTrending = "overallTrend", |
|
|
AbnormalTrend = "abnormalTrend", |
|
|
PearsonCorrelation = "pearsonCorrelation", |
|
|
SpearmanCorrelation = "spearmanCorrelation", |
|
|
ExtremeValue = "extremeValue", |
|
|
MajorityValue = "majorityValue", |
|
|
StatisticsAbnormal = "statisticsAbnormal", |
|
|
StatisticsBase = "statisticsBase", |
|
|
DbscanOutlier = "dbscanOutlier", |
|
|
LOFOutlier = "lofOutlier", |
|
|
TurningPoint = "turningPoint", |
|
|
PageHinkley = "pageHinkley", |
|
|
DifferenceOutlier = "differenceOutlier", |
|
|
Volatility = "volatility", |
|
|
} |
|
|
|
|
|
const getBase64 = async (spec: any, width?: number, height?: number) => { |
|
|
spec.animation = false; |
|
|
width && (spec.width = width); |
|
|
height && (spec.height = height); |
|
|
const browser = await puppeteer.launch(); |
|
|
const page = await browser.newPage(); |
|
|
await page.setContent(getHtmlVChart(spec, width, height)); |
|
|
|
|
|
const dataUrl = await page.evaluate(() => { |
|
|
const canvas: any = document |
|
|
.getElementById("chart-container") |
|
|
?.querySelector("canvas"); |
|
|
return canvas?.toDataURL("image/png"); |
|
|
}); |
|
|
|
|
|
const base64Data = dataUrl.replace(/^data:image\/png;base64,/, ""); |
|
|
await browser.close(); |
|
|
return Buffer.from(base64Data, "base64"); |
|
|
}; |
|
|
|
|
|
const serializeSpec = (spec: any) => { |
|
|
return JSON.stringify(spec, (key, value) => { |
|
|
if (typeof value === "function") { |
|
|
const funcStr = value |
|
|
.toString() |
|
|
.replace(/(\r\n|\n|\r)/gm, "") |
|
|
.replace(/\s+/g, " "); |
|
|
|
|
|
return `__FUNCTION__${funcStr}`; |
|
|
} |
|
|
return value; |
|
|
}); |
|
|
}; |
|
|
|
|
|
function getHtmlVChart(spec: any, width?: number, height?: number) { |
|
|
return `<!DOCTYPE html> |
|
|
<html> |
|
|
<head> |
|
|
<title>VChart Demo</title> |
|
|
<script src="https://unpkg.com/@visactor/vchart/build/index.min.js"></script> |
|
|
</head> |
|
|
<body> |
|
|
<div id="chart-container" style="width: ${ |
|
|
width ? `${width}px` : "100%" |
|
|
}; height: ${height ? `${height}px` : "100%"};"></div> |
|
|
<script> |
|
|
// parse spec with function |
|
|
function parseSpec(stringSpec) { |
|
|
return JSON.parse(stringSpec, (k, v) => { |
|
|
if (typeof v === 'string' && v.startsWith('__FUNCTION__')) { |
|
|
const funcBody = v.slice(12); // 移除标记 |
|
|
try { |
|
|
return new Function('return (' + funcBody + ')')(); |
|
|
} catch(e) { |
|
|
console.error('函数解析失败:', e); |
|
|
return () => {}; |
|
|
} |
|
|
} |
|
|
return v; |
|
|
}); |
|
|
} |
|
|
const spec = parseSpec(\`${serializeSpec(spec)}\`); |
|
|
const chart = new VChart.VChart(spec, { |
|
|
dom: 'chart-container' |
|
|
}); |
|
|
chart.renderSync(); |
|
|
</script> |
|
|
</body> |
|
|
</html> |
|
|
`; |
|
|
} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
function getSavedPathName( |
|
|
directory: string, |
|
|
fileName: string, |
|
|
outputType: "html" | "png" | "json" | "md", |
|
|
isUpdate: boolean = false |
|
|
) { |
|
|
let newFileName = fileName; |
|
|
while ( |
|
|
!isUpdate && |
|
|
fs.existsSync( |
|
|
path.join(directory, "visualization", `${newFileName}.${outputType}`) |
|
|
) |
|
|
) { |
|
|
newFileName += "_new"; |
|
|
} |
|
|
return path.join(directory, "visualization", `${newFileName}.${outputType}`); |
|
|
} |
|
|
|
|
|
const readStdin = (): Promise<string> => { |
|
|
return new Promise((resolve) => { |
|
|
let input = ""; |
|
|
process.stdin.setEncoding("utf-8"); |
|
|
process.stdin.on("data", (chunk) => (input += chunk)); |
|
|
process.stdin.on("end", () => resolve(input)); |
|
|
}); |
|
|
}; |
|
|
|
|
|
|
|
|
const setInsightTemplate = ( |
|
|
path: string, |
|
|
title: string, |
|
|
insights: string[] |
|
|
) => { |
|
|
let res = ""; |
|
|
if (insights.length) { |
|
|
res += `## ${title} Insights`; |
|
|
insights.forEach((insight, index) => { |
|
|
res += `\n${index + 1}. ${insight}`; |
|
|
}); |
|
|
} |
|
|
if (res) { |
|
|
fs.writeFileSync(path, res, "utf-8"); |
|
|
return { insight_path: path, insight_md: res }; |
|
|
} |
|
|
return {}; |
|
|
}; |
|
|
|
|
|
|
|
|
async function saveChartRes(options: { |
|
|
spec: any; |
|
|
directory: string; |
|
|
outputType: "png" | "html"; |
|
|
fileName: string; |
|
|
width?: number; |
|
|
height?: number; |
|
|
isUpdate?: boolean; |
|
|
}) { |
|
|
const { directory, fileName, spec, outputType, width, height, isUpdate } = |
|
|
options; |
|
|
const specPath = getSavedPathName(directory, fileName, "json", isUpdate); |
|
|
fs.writeFileSync(specPath, JSON.stringify(spec, null, 2)); |
|
|
const savedPath = getSavedPathName(directory, fileName, outputType, isUpdate); |
|
|
if (outputType === "png") { |
|
|
const base64 = await getBase64(spec, width, height); |
|
|
fs.writeFileSync(savedPath, base64); |
|
|
} else { |
|
|
const html = getHtmlVChart(spec, width, height); |
|
|
fs.writeFileSync(savedPath, html, "utf-8"); |
|
|
} |
|
|
return savedPath; |
|
|
} |
|
|
|
|
|
async function generateChart( |
|
|
vmind: VMind, |
|
|
options: { |
|
|
dataset: string | DataTable; |
|
|
userPrompt: string; |
|
|
directory: string; |
|
|
outputType: "png" | "html"; |
|
|
fileName: string; |
|
|
width?: number; |
|
|
height?: number; |
|
|
language?: "en" | "zh"; |
|
|
} |
|
|
) { |
|
|
let res: { |
|
|
chart_path?: string; |
|
|
error?: string; |
|
|
insight_path?: string; |
|
|
insight_md?: string; |
|
|
} = {}; |
|
|
const { |
|
|
dataset, |
|
|
userPrompt, |
|
|
directory, |
|
|
width, |
|
|
height, |
|
|
outputType, |
|
|
fileName, |
|
|
language, |
|
|
} = options; |
|
|
try { |
|
|
|
|
|
const jsonDataset = isString(dataset) ? JSON.parse(dataset) : dataset; |
|
|
const { spec, error, chartType } = await vmind.generateChart( |
|
|
userPrompt, |
|
|
undefined, |
|
|
jsonDataset, |
|
|
{ |
|
|
enableDataQuery: false, |
|
|
theme: "light", |
|
|
} |
|
|
); |
|
|
if (error || !spec) { |
|
|
return { |
|
|
error: error || "Spec of Chart was Empty!", |
|
|
}; |
|
|
} |
|
|
|
|
|
spec.title = { |
|
|
text: userPrompt, |
|
|
}; |
|
|
if (!fs.existsSync(path.join(directory, "visualization"))) { |
|
|
fs.mkdirSync(path.join(directory, "visualization")); |
|
|
} |
|
|
const specPath = getSavedPathName(directory, fileName, "json"); |
|
|
res.chart_path = await saveChartRes({ |
|
|
directory, |
|
|
spec, |
|
|
width, |
|
|
height, |
|
|
fileName, |
|
|
outputType, |
|
|
}); |
|
|
|
|
|
|
|
|
const insights = []; |
|
|
if ( |
|
|
chartType && |
|
|
[ |
|
|
ChartType.BarChart, |
|
|
ChartType.LineChart, |
|
|
ChartType.AreaChart, |
|
|
ChartType.ScatterPlot, |
|
|
ChartType.DualAxisChart, |
|
|
].includes(chartType) |
|
|
) { |
|
|
const { insights: vmindInsights } = await vmind.getInsights(spec, { |
|
|
maxNum: 6, |
|
|
algorithms: [ |
|
|
AlgorithmType.OverallTrending, |
|
|
AlgorithmType.AbnormalTrend, |
|
|
AlgorithmType.PearsonCorrelation, |
|
|
AlgorithmType.SpearmanCorrelation, |
|
|
AlgorithmType.StatisticsAbnormal, |
|
|
AlgorithmType.LOFOutlier, |
|
|
AlgorithmType.DbscanOutlier, |
|
|
AlgorithmType.MajorityValue, |
|
|
AlgorithmType.PageHinkley, |
|
|
AlgorithmType.TurningPoint, |
|
|
AlgorithmType.StatisticsBase, |
|
|
AlgorithmType.Volatility, |
|
|
], |
|
|
usePolish: false, |
|
|
language: language === "en" ? "english" : "chinese", |
|
|
}); |
|
|
insights.push(...vmindInsights); |
|
|
} |
|
|
const insightsText = insights |
|
|
.map((insight) => insight.textContent?.plainText) |
|
|
.filter((insight) => !!insight) as string[]; |
|
|
spec.insights = insights; |
|
|
fs.writeFileSync(specPath, JSON.stringify(spec, null, 2)); |
|
|
res = { |
|
|
...res, |
|
|
...setInsightTemplate( |
|
|
getSavedPathName(directory, fileName, "md"), |
|
|
userPrompt, |
|
|
insightsText |
|
|
), |
|
|
}; |
|
|
} catch (error: any) { |
|
|
res.error = error.toString(); |
|
|
} finally { |
|
|
return res; |
|
|
} |
|
|
} |
|
|
|
|
|
async function updateChartWithInsight( |
|
|
vmind: VMind, |
|
|
options: { |
|
|
directory: string; |
|
|
outputType: "png" | "html"; |
|
|
fileName: string; |
|
|
insightsId: number[]; |
|
|
} |
|
|
) { |
|
|
const { directory, outputType, fileName, insightsId } = options; |
|
|
let res: { error?: string; chart_path?: string } = {}; |
|
|
try { |
|
|
const specPath = getSavedPathName(directory, fileName, "json", true); |
|
|
const spec = JSON.parse(fs.readFileSync(specPath, "utf8")); |
|
|
|
|
|
const insights = (spec.insights || []).filter( |
|
|
(_insight: any, index: number) => insightsId.includes(index + 1) |
|
|
); |
|
|
const { newSpec, error } = await vmind.updateSpecByInsights(spec, insights); |
|
|
if (error) { |
|
|
throw error; |
|
|
} |
|
|
res.chart_path = await saveChartRes({ |
|
|
spec: newSpec, |
|
|
directory, |
|
|
outputType, |
|
|
fileName, |
|
|
isUpdate: true, |
|
|
}); |
|
|
} catch (error: any) { |
|
|
res.error = error.toString(); |
|
|
} finally { |
|
|
return res; |
|
|
} |
|
|
} |
|
|
|
|
|
async function executeVMind() { |
|
|
const input = await readStdin(); |
|
|
const inputData = JSON.parse(input); |
|
|
let res; |
|
|
const { |
|
|
llm_config, |
|
|
width, |
|
|
dataset = [], |
|
|
height, |
|
|
directory, |
|
|
user_prompt: userPrompt, |
|
|
output_type: outputType = "png", |
|
|
file_name: fileName, |
|
|
task_type: taskType = "visualization", |
|
|
insights_id: insightsId = [], |
|
|
language = "en", |
|
|
} = inputData; |
|
|
const { base_url: baseUrl, model, api_key: apiKey } = llm_config; |
|
|
const vmind = new VMind({ |
|
|
url: `${baseUrl}/chat/completions`, |
|
|
model, |
|
|
headers: { |
|
|
"api-key": apiKey, |
|
|
Authorization: `Bearer ${apiKey}`, |
|
|
}, |
|
|
}); |
|
|
if (taskType === "visualization") { |
|
|
res = await generateChart(vmind, { |
|
|
dataset, |
|
|
userPrompt, |
|
|
directory, |
|
|
outputType, |
|
|
fileName, |
|
|
width, |
|
|
height, |
|
|
language, |
|
|
}); |
|
|
} else if (taskType === "insight" && insightsId.length) { |
|
|
res = await updateChartWithInsight(vmind, { |
|
|
directory, |
|
|
fileName, |
|
|
outputType, |
|
|
insightsId, |
|
|
}); |
|
|
} |
|
|
console.log(JSON.stringify(res)); |
|
|
} |
|
|
|
|
|
executeVMind(); |
|
|
|
