@visactor/vmind/esm/atom/chartGenerator/index.js

<div align="center"> <a href="https://github.com/VisActor#gh-light-mode-only" target="_blank"> <img alt="VisActor Logo" width="200" src="https://github.com/VisActor/.github/blob/main/profile/logo_500_200_light.svg"/> </a> <a href="https://githu

import { AtomName } from "../../types/atom";

import { BaseAtom } from "../base";

import { merge } from "@visactor/vutils";

import { DEFAULT_MAP_OPTION, SUPPORTED_CHART_LIST } from "./const";

import { getPrompt, revisedUserInput } from "./prompt";

import { getContextAfterRevised } from "./llmResultRevise";

import { checkChartTypeAndCell, getVizSchema } from "./utils";

import { getChartSpecWithContext } from "./spec";

import { getCellContextBySimpleVChartSpec, getRuleLLMContent } from "./rule";

import { getCellContextByAdvisor } from "./advisor";

import { getFieldInfoFromDataset } from "../../utils/field";

import { Factory } from "../../core/factory";

export class ChartGeneratorAtom extends BaseAtom {
    constructor(context, option) {
        super(context, option), this.name = AtomName.CHART_GENERATE, this.isLLMAtom = !0, 
        this._generateType = "llm", this.setFinalChartTypeList();
    }
    setFinalChartTypeList() {
        const {chartTypeList: chartTypeList, unsupportChartTypeList: unsupportChartTypeList} = this.options;
        this.finalChartTypeList = chartTypeList.filter((v => !unsupportChartTypeList.includes(v)));
    }
    buildDefaultContext(context) {
        return merge({}, {
            dataTable: [],
            fieldInfo: []
        }, context);
    }
    buildDefaultOptions() {
        return Object.assign(Object.assign({}, super.buildDefaultOptions()), {
            useChartAdvisor: !1,
            chartTypeList: SUPPORTED_CHART_LIST,
            basemapOption: DEFAULT_MAP_OPTION,
            unsupportChartTypeList: [],
            useChartRule: !1
        });
    }
    updateContext(context) {
        return this.context = super.updateContext(context), this.context.vizSchema = getVizSchema(this.context), 
        this.context;
    }
    updateOptions(options) {
        super.updateOptions(options), this.setFinalChartTypeList();
    }
    getLLMMessages(query) {
        const {command: command} = this.context, {showThoughts: showThoughts} = this.options, addtionContent = this.getHistoryLLMMessages(query);
        return [ {
            role: "system",
            content: getPrompt(this.finalChartTypeList, showThoughts)
        }, {
            role: "user",
            content: revisedUserInput(command, this.context.vizSchema)
        }, ...addtionContent ];
    }
    parseLLMContent(resJson) {
        const {CHART_TYPE: CHART_TYPE, FIELD_MAP: FIELD_MAP, thoughts: thoughts, stackOrPercent: stackOrPercent, transpose: transpose} = resJson;
        let newContext = Object.assign(Object.assign({}, this.context), {
            thoughts: thoughts,
            chartType: CHART_TYPE,
            cell: FIELD_MAP,
            chartTypeList: this.finalChartTypeList,
            stackOrPercent: stackOrPercent,
            transpose: transpose
        });
        newContext = getContextAfterRevised(newContext);
        const {error: error, chartType: chartType, fieldInfo: fieldInfo, cell: cell} = newContext;
        return !error && checkChartTypeAndCell(chartType, cell, fieldInfo) || (console.warn("LLM generation error, use rule generation."), 
        delete newContext.error, delete newContext.message), newContext;
    }
    runBeforeLLM() {
        const {dataTable: dataTable, fieldInfo: fieldInfo, simpleVChartSpec: simpleVChartSpec} = this.context;
        if (simpleVChartSpec) {
            this.isLLMAtom = !1, this._generateType = "simpleSpec";
            const {ctx: ctx, mockLLMContent: mockLLMContent} = getCellContextBySimpleVChartSpec(simpleVChartSpec);
            this.updateContext(ctx), this.updateContext(this.parseLLMContent(mockLLMContent));
        } else {
            if (fieldInfo && 0 !== fieldInfo.length || this.updateContext({
                fieldInfo: getFieldInfoFromDataset(dataTable)
            }), (dataTable.length > 1 || !this.options.useChartRule) && !simpleVChartSpec) return this._generateType = this.options.useChartAdvisor ? "chartAdvistor" : "llm", 
            this.isLLMAtom = "llm" === this._generateType, this.context;
            this.isLLMAtom = !1, this._generateType = "rule";
            const ruleResJson = getRuleLLMContent(this.context);
            ruleResJson ? this.updateContext(this.parseLLMContent(ruleResJson)) : this.updateContext({
                cell: null
            });
        }
        return this.context;
    }
    runWithLLMError(error) {
        return super._runWithOutLLM(), this._generateType = "chartAdvistor", this._runWithOutLLM();
    }
    _runWithOutLLM() {
        var _a, _b, _c, _d;
        if (this.isLLMAtom = !0, "rule" === this._generateType && !this.context.cell) return this.context.spec = null, 
        this.context;
        const additionalCtx = {
            chartTypeList: this.finalChartTypeList,
            basemapOption: null === (_a = this.options) || void 0 === _a ? void 0 : _a.basemapOption,
            totalTime: null === (_b = this.options) || void 0 === _b ? void 0 : _b.animationDuration,
            colors: null === (_c = this.options) || void 0 === _c ? void 0 : _c.colorPalette,
            chartTheme: null === (_d = this.options) || void 0 === _d ? void 0 : _d.theme
        };
        if ("chartAdvistor" === this._generateType || this.options.useChartAdvisor) {
            const {cell: cell, dataset: dataset, chartType: chartType, advisedList: advisedList, usage: usage} = getCellContextByAdvisor(Object.assign(Object.assign({}, this.context), additionalCtx));
            this.context = Object.assign(Object.assign({}, this.context), {
                usage: usage,
                cell: cell,
                dataTable: dataset,
                chartType: chartType,
                chartAdvistorRes: advisedList.map((item => {
                    const tmpContext = merge({}, this.context, {
                        cell: item.cell,
                        dataTable: item.dataset,
                        chartType: item.chartType
                    });
                    return {
                        chartType: item.chartType,
                        score: item.score,
                        spec: getChartSpecWithContext(Object.assign(Object.assign({}, tmpContext), additionalCtx)).spec
                    };
                }))
            });
        } else this.context.chartAdvistorRes = [];
        const newContext = Object.assign(Object.assign({}, this.context), getChartSpecWithContext(Object.assign(Object.assign({}, this.context), additionalCtx)));
        return this.updateContext(newContext), newContext;
    }
}

export const registerChartGeneratorAtom = () => {
    Factory.registerAtom(AtomName.CHART_GENERATE, ChartGeneratorAtom);
};
//# sourceMappingURL=index.js.map