1
export class WebLlmVectorProvider {
2
    /** @type {object?} WebLLM engine */
3
    #engine = null;
4

5
    constructor() {
6
        this.#engine = null;
7
    }
8

9
    /**
10
     * Check if WebLLM is available and up-to-date
11
     * @throws {Error} If WebLLM is not available or not up-to-date
12
     */
13
    #checkWebLlm() {
14
        if (!Object.hasOwn(SillyTavern, 'llm')) {
15
            throw new Error('WebLLM is not available', { cause: 'webllm-not-available' });
16
        }
17

18
        if (typeof SillyTavern.llm.generateEmbedding !== 'function') {
19
            throw new Error('WebLLM is not updated', { cause: 'webllm-not-updated' });
20
        }
21
    }
22

23
    /**
24
     * Initialize the engine with a model.
25
     * @param {string} modelId Model ID to initialize the engine with
26
     * @returns {Promise<void>} Promise that resolves when the engine is initialized
27
     */
28
    #initEngine(modelId) {
29
        this.#checkWebLlm();
30
        if (!this.#engine) {
31
            this.#engine = SillyTavern.llm.getEngine();
32
        }
33

34
        return this.#engine.loadModel(modelId);
35
    }
36

37
    /**
38
     * Get available models.
39
     * @returns {{id:string, toString: function(): string}[]} Array of available models
40
     */
41
    getModels() {
42
        this.#checkWebLlm();
43
        return SillyTavern.llm.getEmbeddingModels();
44
    }
45

46
    /**
47
     * Generate embeddings for a list of texts.
48
     * @param {string[]} texts Array of texts to generate embeddings for
49
     * @param {string} modelId Model to use for generating embeddings
50
     * @returns {Promise<number[][]>} Array of embeddings for each text
51
     */
52
    async embedTexts(texts, modelId) {
53
        await this.#initEngine(modelId);
54
        return this.#engine.generateEmbedding(texts);
55
    }
56

57
    /**
58
     * Loads a model into the engine.
59
     * @param {string} modelId Model ID to load
60
     */
61
    async loadModel(modelId) {
62
        await this.#initEngine(modelId);
63
    }
64
}