generator-begcode/dist/agent/agent-core/rag/StandardRagBuilder.js

Spring Boot + Angular/React/Vue in one handy generator

import { v4 as uuid } from 'uuid';
import { InMemoryWorkspace } from '../../agent-utils/index.js';
import { LocalVectorDB } from '../embeddings/index.js';
export class QueryWithRecombine {
    queryOrVector;
    collection;
    _items;
    addToCollectionPromise;
    constructor(queryOrVector, collection, _items, addToCollectionPromise) {
        this.queryOrVector = queryOrVector;
        this.collection = collection;
        this._items = _items;
        this.addToCollectionPromise = addToCollectionPromise;
    }
    recombine(recombiner) {
        return recombiner(async () => {
            await this.addToCollectionPromise();
            return this.collection.iterativeSearch(this.queryOrVector);
        }, this._items);
    }
}
export class StandardRagBuilder {
    _selector = x => x;
    collection;
    _items;
    lastAddedItemIndex;
    constructor(context, collectionName, workspace, items) {
        const db = new LocalVectorDB(workspace ?? new InMemoryWorkspace(), 'ragdb', context.embedding);
        this.collection = db.addCollection(collectionName ?? uuid());
        this._items = items ?? [];
        this.lastAddedItemIndex = items ? items.length - 1 : -1;
    }
    addItems(items) {
        this._items.push(...items);
        return this;
    }
    async forceAddItemsToCollection() {
        await this.addItemsToCollection();
        return this;
    }
    selector(selector) {
        this._selector = selector;
        return this;
    }
    query(queryOrVector) {
        return new QueryWithRecombine(queryOrVector, this.collection, this._items, this.forceAddItemsToCollection.bind(this));
    }
    async addItemsToCollection() {
        const itemsToAdd = this._items.slice(this.lastAddedItemIndex + 1);
        if (!itemsToAdd.length) {
            return;
        }
        await this.collection.add(itemsToAdd.map(x => {
            const text = this._selector(x);
            if (typeof text != 'string') {
                throw Error('Selector should return a string. Perhaps you forgot to set the selector?');
            }
            return text;
        }), itemsToAdd.map((_, i) => ({ index: this.lastAddedItemIndex + i + 1 })));
        this.lastAddedItemIndex = this._items.length - 1;
    }
}