@visactor/vmind/esm/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

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

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

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

import { getQueryDatasetPrompt } from "./prompt";

import { parseSQLResponse } from "./utils";

import { executeDataQuery, getFinalQueryResult, patchSQLBeforeQuery, restoreDatasetAfterQuery } from "./executeQuery";

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

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

export class DataQueryAtom extends BaseAtom {
    constructor(context, option) {
        super(context, option), this.name = AtomName.DATA_QUERY, this.isLLMAtom = !0;
    }
    buildDefaultContext(context) {
        return 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 => pick(info, [ "fieldName", "type", "role" ])));
            return [ {
                role: "system",
                content: 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), 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 = getFieldInfoFromDataset(dataTable)), 
        this.context;
    }
    _runWithOutLLM() {
        let newContext = Object.assign({}, this.context);
        return [ patchSQLBeforeQuery, executeDataQuery, restoreDatasetAfterQuery, 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;
    }
}

export const registerDataQueryAtom = () => {
    Factory.registerAtom(AtomName.DATA_QUERY, DataQueryAtom);
};
//# sourceMappingURL=index.js.map