@llamaindex/core/retriever/dist/index.cjs

LlamaIndex Core Module

Object.defineProperty(exports, '__esModule', { value: true });

var env = require('@llamaindex/env');
var index_cjs$1 = require('../../global/dist/index.cjs');
var index_cjs = require('../../prompts/dist/index.cjs');
var index_cjs$2 = require('../../schema/dist/index.cjs');

class BaseRetriever extends index_cjs.PromptMixin {
    _updatePrompts() {}
    _getPrompts() {
        return {};
    }
    _getPromptModules() {
        return {};
    }
    constructor(){
        super(), this.objectMap = new Map();
    }
    async retrieve(params) {
        const cb = index_cjs$1.Settings.callbackManager;
        const queryBundle = typeof params === "string" ? {
            query: params
        } : params;
        const id = env.randomUUID();
        cb.dispatchEvent("retrieve-start", {
            id,
            query: queryBundle
        });
        let response = await this._retrieve(queryBundle);
        response = await this._handleRecursiveRetrieval(queryBundle, response);
        cb.dispatchEvent("retrieve-end", {
            id,
            query: queryBundle,
            nodes: response
        });
        return response;
    }
    async _handleRecursiveRetrieval(params, nodes) {
        const retrievedNodes = [];
        for (const { node, score = 1.0 } of nodes){
            if (node.type === index_cjs$2.ObjectType.INDEX) {
                const indexNode = node;
                const object = this.objectMap.get(indexNode.indexId);
                if (object !== undefined) {
                    retrievedNodes.push(...this._retrieveFromObject(object, params, score));
                } else {
                    retrievedNodes.push({
                        node,
                        score
                    });
                }
            } else {
                retrievedNodes.push({
                    node,
                    score
                });
            }
        }
        return nodes;
    }
    _retrieveFromObject(object, queryBundle, score) {
        if (object == null) {
            throw new TypeError("Object is not retrievable");
        }
        if (typeof object !== "object") {
            throw new TypeError("Object is not retrievable");
        }
        if ("node" in object && object.node instanceof index_cjs$2.BaseNode) {
            return [
                {
                    node: object.node,
                    score: "score" in object && typeof object.score === "number" ? object.score : score
                }
            ];
        }
        if (object instanceof index_cjs$2.BaseNode) {
            return [
                {
                    node: object,
                    score
                }
            ];
        } else {
            // todo: support other types
            // BaseQueryEngine
            // BaseRetriever
            // QueryComponent
            throw new TypeError("Object is not retrievable");
        }
    }
}

exports.BaseRetriever = BaseRetriever;