File size: 1,900 Bytes
519a20c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
export class WebLlmVectorProvider {
    /** @type {object?} WebLLM engine */
    #engine = null;

    constructor() {
        this.#engine = null;
    }

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

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

    /**
     * Initialize the engine with a model.
     * @param {string} modelId Model ID to initialize the engine with
     * @returns {Promise<void>} Promise that resolves when the engine is initialized
     */
    #initEngine(modelId) {
        this.#checkWebLlm();
        if (!this.#engine) {
            this.#engine = SillyTavern.llm.getEngine();
        }

        return this.#engine.loadModel(modelId);
    }

    /**
     * Get available models.
     * @returns {{id:string, toString: function(): string}[]} Array of available models
     */
    getModels() {
        this.#checkWebLlm();
        return SillyTavern.llm.getEmbeddingModels();
    }

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

    /**
     * Loads a model into the engine.
     * @param {string} modelId Model ID to load
     */
    async loadModel(modelId) {
        await this.#initEngine(modelId);
    }
}