@visactor/vmind/cjs/atom/dataQuery/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

"use strict";

Object.defineProperty(exports, "__esModule", {
    value: !0
}), exports.registerDataQueryAtom = exports.DataQueryAtom = void 0;

const atom_1 = require("../../types/atom"), base_1 = require("../base"), vutils_1 = require("@visactor/vutils"), prompt_1 = require("./prompt"), utils_1 = require("./utils"), executeQuery_1 = require("./executeQuery"), field_1 = require("../../utils/field"), factory_1 = require("../../core/factory");

class DataQueryAtom extends base_1.BaseAtom {
    constructor(context, option) {
        super(context, option), this.name = atom_1.AtomName.DATA_QUERY, this.isLLMAtom = !0;
    }
    buildDefaultContext(context) {
        return (0, vutils_1.merge)({}, {
            dataTable: [],
            fieldInfo: [],
            llmFieldInfo: [],
            command: "",
            dataTableSummary: ""
        }, context);
    }
    buildDefaultOptions() {
        return Object.assign(Object.assign({}, super.buildDefaultOptions()), {
            useSQL: !0
        });
    }
    getLLMMessages(query) {
        const {fieldInfo: fieldInfo, command: command} = this.context, {showThoughts: showThoughts} = this.options, addtionContent = this.getHistoryLLMMessages(query);
        if (this.options.useSQL) {
            const fieldInfoContent = fieldInfo.map((info => (0, vutils_1.pick)(info, [ "fieldName", "type", "role" ])));
            return [ {
                role: "system",
                content: (0, prompt_1.getQueryDatasetPrompt)(showThoughts)
            }, {
                role: "user",
                content: `User's Command: ${command}\nColumn Information: ${JSON.stringify(fieldInfoContent)}`
            }, ...addtionContent ];
        }
        return [];
    }
    parseLLMContent(resJson, toolJson, llmRes) {
        var _a;
        const {sql: sql, fieldInfo: responseFiledInfo, thoughts: thoughts = ""} = resJson;
        if ((!sql || !responseFiledInfo) && (null === (_a = null == llmRes ? void 0 : llmRes.choices) || void 0 === _a ? void 0 : _a[0])) {
            const content = llmRes.choices[0].message.content;
            return Object.assign(Object.assign({}, this.context), (0, utils_1.parseSQLResponse)(content));
        }
        return Object.assign(Object.assign({}, this.context), {
            sql: sql,
            llmFieldInfo: responseFiledInfo,
            thoughts: thoughts
        });
    }
    runBeforeLLM() {
        const {fieldInfo: fieldInfo = [], dataTable: dataTable} = this.context;
        return !fieldInfo.length && dataTable.length && (this.context.fieldInfo = (0, field_1.getFieldInfoFromDataset)(dataTable)), 
        this.context;
    }
    _runWithOutLLM() {
        let newContext = Object.assign({}, this.context);
        return [ executeQuery_1.patchSQLBeforeQuery, executeQuery_1.executeDataQuery, executeQuery_1.restoreDatasetAfterQuery, executeQuery_1.getFinalQueryResult ].forEach((func => {
            newContext = Object.assign(Object.assign({}, newContext), func(newContext));
        })), this.setNewContext(Object.assign(Object.assign({}, this.context), {
            dataTable: newContext.dataTable,
            fieldInfo: newContext.fieldInfo
        })), this.context;
    }
}

exports.DataQueryAtom = DataQueryAtom;

const registerDataQueryAtom = () => {
    factory_1.Factory.registerAtom(atom_1.AtomName.DATA_QUERY, DataQueryAtom);
};

exports.registerDataQueryAtom = registerDataQueryAtom;
//# sourceMappingURL=index.js.map