"use strict";
Object.defineProperty(exports, "__esModule", { value: true });
const lodash_1 = require("lodash");
const llamaindex_1 = require("llamaindex");
const EngineUtils_1 = require("../EngineUtils");
const EvaluationRunTracerLlama_1 = require("../../../evaluation/EvaluationRunTracerLlama");
class SubQuestionQueryEngine_LlamaIndex {
    constructor(fields) {
        this.label = 'Sub Question Query Engine';
        this.name = 'subQuestionQueryEngine';
        this.version = 2.0;
        this.type = 'SubQuestionQueryEngine';
        this.icon = 'subQueryEngine.svg';
        this.category = 'Engine';
        this.description =
            'Breaks complex query into sub questions for each relevant data source, then gather all the intermediate responses and synthesizes a final response';
        this.baseClasses = [this.type, 'BaseQueryEngine'];
        this.tags = ['LlamaIndex'];
        this.badge = 'DEPRECATING';
        this.deprecateMessage = 'LlamaIndex integration is deprecated and will be removed in a future release.';
        this.inputs = [
            {
                label: 'QueryEngine Tools',
                name: 'queryEngineTools',
                type: 'QueryEngineTool',
                list: true
            },
            {
                label: 'Chat Model',
                name: 'model',
                type: 'BaseChatModel_LlamaIndex'
            },
            {
                label: 'Embeddings',
                name: 'embeddings',
                type: 'BaseEmbedding_LlamaIndex'
            },
            {
                label: 'Response Synthesizer',
                name: 'responseSynthesizer',
                type: 'ResponseSynthesizer',
                description: 'ResponseSynthesizer is responsible for sending the query, nodes, and prompt templates to the LLM to generate a response. See <a target="_blank" href="https://ts.llamaindex.ai/modules/response_synthesizer">more</a>',
                optional: true
            },
            {
                label: 'Return Source Documents',
                name: 'returnSourceDocuments',
                type: 'boolean',
                optional: true
            }
        ];
        this.sessionId = fields?.sessionId;
    }
    async init(nodeData) {
        return prepareEngine(nodeData);
    }
    async run(nodeData, input, options) {
        const returnSourceDocuments = nodeData.inputs?.returnSourceDocuments;
        const queryEngine = prepareEngine(nodeData);
        let text = '';
        let sourceDocuments = [];
        let sourceNodes = [];
        let isStreamingStarted = false;
        await EvaluationRunTracerLlama_1.EvaluationRunTracerLlama.injectEvaluationMetadata(nodeData, options, queryEngine);
        const shouldStreamResponse = options.shouldStreamResponse;
        const sseStreamer = options.sseStreamer;
        const chatId = options.chatId;
        if (shouldStreamResponse) {
            const stream = await queryEngine.query({ query: input, stream: true });
            for await (const chunk of stream) {
                text += chunk.response;
                if (chunk.sourceNodes)
                    sourceNodes = chunk.sourceNodes;
                if (!isStreamingStarted) {
                    isStreamingStarted = true;
                    if (sseStreamer) {
                        sseStreamer.streamStartEvent(chatId, chunk.response);
                    }
                }
                if (sseStreamer) {
                    sseStreamer.streamTokenEvent(chatId, chunk.response);
                }
            }
            if (returnSourceDocuments) {
                sourceDocuments = (0, EngineUtils_1.reformatSourceDocuments)(sourceNodes);
                if (sseStreamer) {
                    sseStreamer.streamSourceDocumentsEvent(chatId, sourceDocuments);
                }
            }
        }
        else {
            const response = await queryEngine.query({ query: input });
            text = response?.response;
            sourceDocuments = (0, EngineUtils_1.reformatSourceDocuments)(response?.sourceNodes ?? []);
        }
        if (returnSourceDocuments)
            return { text, sourceDocuments };
        else
            return { text };
    }
}
const prepareEngine = (nodeData) => {
    const embeddings = nodeData.inputs?.embeddings;
    const model = nodeData.inputs?.model;
    const serviceContext = (0, llamaindex_1.serviceContextFromDefaults)({
        llm: model,
        embedModel: embeddings
    });
    let queryEngineTools = nodeData.inputs?.queryEngineTools;
    queryEngineTools = (0, lodash_1.flatten)(queryEngineTools);
    let queryEngine = llamaindex_1.SubQuestionQueryEngine.fromDefaults({
        serviceContext,
        queryEngineTools,
        questionGen: new llamaindex_1.LLMQuestionGenerator({ llm: model })
    });
    const responseSynthesizerObj = nodeData.inputs?.responseSynthesizer;
    if (responseSynthesizerObj) {
        if (responseSynthesizerObj.type === 'TreeSummarize') {
            const responseSynthesizer = new llamaindex_1.ResponseSynthesizer({
                responseBuilder: new llamaindex_1.TreeSummarize(serviceContext, responseSynthesizerObj.textQAPromptTemplate),
                serviceContext
            });
            queryEngine = llamaindex_1.SubQuestionQueryEngine.fromDefaults({
                responseSynthesizer,
                serviceContext,
                queryEngineTools,
                questionGen: new llamaindex_1.LLMQuestionGenerator({ llm: model })
            });
        }
        else if (responseSynthesizerObj.type === 'CompactAndRefine') {
            const responseSynthesizer = new llamaindex_1.ResponseSynthesizer({
                responseBuilder: new llamaindex_1.CompactAndRefine(serviceContext, responseSynthesizerObj.textQAPromptTemplate, responseSynthesizerObj.refinePromptTemplate),
                serviceContext
            });
            queryEngine = llamaindex_1.SubQuestionQueryEngine.fromDefaults({
                responseSynthesizer,
                serviceContext,
                queryEngineTools,
                questionGen: new llamaindex_1.LLMQuestionGenerator({ llm: model })
            });
        }
        else if (responseSynthesizerObj.type === 'Refine') {
            const responseSynthesizer = new llamaindex_1.ResponseSynthesizer({
                responseBuilder: new llamaindex_1.Refine(serviceContext, responseSynthesizerObj.textQAPromptTemplate, responseSynthesizerObj.refinePromptTemplate),
                serviceContext
            });
            queryEngine = llamaindex_1.SubQuestionQueryEngine.fromDefaults({
                responseSynthesizer,
                serviceContext,
                queryEngineTools,
                questionGen: new llamaindex_1.LLMQuestionGenerator({ llm: model })
            });
        }
        else if (responseSynthesizerObj.type === 'SimpleResponseBuilder') {
            const responseSynthesizer = new llamaindex_1.ResponseSynthesizer({
                responseBuilder: new llamaindex_1.SimpleResponseBuilder(serviceContext),
                serviceContext
            });
            queryEngine = llamaindex_1.SubQuestionQueryEngine.fromDefaults({
                responseSynthesizer,
                serviceContext,
                queryEngineTools,
                questionGen: new llamaindex_1.LLMQuestionGenerator({ llm: model })
            });
        }
    }
    return queryEngine;
};
module.exports = { nodeClass: SubQuestionQueryEngine_LlamaIndex };
//# sourceMappingURL=SubQuestionQueryEngine.js.map