Add support for KoboldCpp embeddings in Vector Storage (#3795) * Add support for KoboldCpp embeddings in vector processing * Add validation for KoboldCpp embeddings to handle empty data * Improve toast handling

80e821d12dbca4827081e388a501a8e8ed4f9527

Cohee <18619528+Cohee1207@users.noreply.github.com>

Signed

Parents: 9c4404cae94db770b798daf6ee7415337028430b

public/scripts/extensions/vectors/index.js+51 -1 Viewed

  * @returns {Promise<boolean>} True if successful, false if not
  */
 async function vectorizeFile(fileText, fileName, collectionId, chunkSize, overlapPercent) {
 +    let toast = jQuery();
++
     try {
         if (settings.translate_files && typeof globalThis.translate === 'function') {
             console.log(`Vectors: Translating file ${fileName} to English...`);

         const batchSize = getBatchSize();
         const toastBody = $('<span>').text('This may take a while. Please wait...');
         const toast = toastr.info(toastBody, `Ingesting file ${escapeHtml(fileName)}`, { closeButton: false, escapeHtml: false, timeOut: 0, extendedTimeOut: 0 });
         const overlapSize = Math.round(chunkSize * overlapPercent / 100);
         const delimiters = getChunkDelimiters();
         // Overlap should not be included in chunk size. It will be later compensated by overlapChunks

         console.log(`Vectors: Inserted ${chunks.length} vector items for file ${fileName} into ${collectionId}`);
         return true;
     } catch (error) {
 +        toastr.clear(toast);
         toastr.error(String(error), 'Failed to vectorize file', { preventDuplicates: true });
         console.error('Vectors: Failed to vectorize file', error);
         return false;

         case 'webllm':
             args.embeddings = await createWebLlmEmbeddings(items);
             break;
 +        case 'koboldcpp': {
 +            const { embeddings, model } = await createKoboldCppEmbeddings(items);
 +            args.embeddings = embeddings;
 +            args.model = model;
 +            break;
 +        }
+    }
     return args;
+}

     if (settings.source === 'ollama' && !textgenerationwebui_settings.server_urls[textgen_types.OLLAMA] ||
         settings.source === 'vllm' && !textgenerationwebui_settings.server_urls[textgen_types.VLLM] ||
 +        settings.source === 'koboldcpp' && !textgenerationwebui_settings.server_urls[textgen_types.KOBOLDCPP] ||
         settings.source === 'llamacpp' && !textgenerationwebui_settings.server_urls[textgen_types.LLAMACPP]) {
         throw new Error('Vectors: API URL missing', { cause: 'api_url_missing' });
+    }

     $('#vllm_vectorsModel').toggle(settings.source === 'vllm');
     $('#nomicai_apiKey').toggle(settings.source === 'nomicai');
     $('#webllm_vectorsModel').toggle(settings.source === 'webllm');
 +    $('#koboldcpp_vectorsModel').toggle(settings.source === 'koboldcpp');
     if (settings.source === 'webllm') {
         loadWebLlmModels();
+    }

     });
+}
 +/**
 + * Creates KoboldCpp embeddings for a list of items.
 + * @param {string[]} items Items to embed
 + * @returns {Promise<{embeddings: Record<string, number[]>, model: string}>} Calculated embeddings
 + */
 +async function createKoboldCppEmbeddings(items) {
 +    const response = await fetch('/api/backends/kobold/embed', {
 +        method: 'POST',
 +        headers: getRequestHeaders(),
 +        body: JSON.stringify({
 +            items: items,
 +            server: textgenerationwebui_settings.server_urls[textgen_types.KOBOLDCPP],
 +        }),
 +    });
++
 +    if (!response.ok) {
 +        throw new Error('Failed to get KoboldCpp embeddings');
 +    }
++
 +    const data = await response.json();
 +    if (!Array.isArray(data.embeddings) || !data.model || data.embeddings.length !== items.length) {
 +        throw new Error('Invalid response from KoboldCpp embeddings');
 +    }
++
 +    const embeddings = /** @type {Record<string, number[]>} */ ({});
 +    for (let i = 0; i < data.embeddings.length; i++) {
 +        if (!Array.isArray(data.embeddings[i]) || data.embeddings[i].length === 0) {
 +            throw new Error('KoboldCpp returned an empty embedding. Reduce the chunk size and/or size threshold and try again.');
 +        }
++
 +        embeddings[items[i]] = data.embeddings[i];
 +    }
++
 +    return {
 +        embeddings: embeddings,
 +        model: data.model,
 +    };
 +}
++
 async function onPurgeClick() {
     const chatId = getCurrentChatId();
     if (!chatId) {

public/scripts/extensions/vectors/settings.html+9 -0 Viewed

                     <option value="cohere">Cohere</option>
                     <option value="extras">Extras (deprecated)</option>
                     <option value="palm">Google AI Studio</option>
 +                    <option value="koboldcpp">KoboldCpp</option>
                     <option value="llamacpp">llama.cpp</option>
                     <option value="transformers" data-i18n="Local (Transformers)">Local (Transformers)</option>
                     <option value="mistral">MistralAI</option>

                     Hint: Set the URL in the API connection settings.
                 </i>
             </div>
 +            <div class="flex-container flexFlowColumn" id="koboldcpp_vectorsModel">
 +                <span>
 +                    Set the KoboldCpp URL in the Text Completion API connection settings.
 +                </span>
 +                <span>
 +                    Must use version 1.87 or higher and have an embedding model loaded.
 +                </span>
 +            </div>
             <div class="flex-container flexFlowColumn" id="llamacpp_vectorsModel">
                 <span data-i18n="The server MUST be started with the --embedding flag to use this feature!">
                     The server MUST be started with the <code>--embedding</code> flag to use this feature!

src/endpoints/backends/kobold.js+42 -0 Viewed

         response.status(500).send('Internal server error');
+    }
 });
++
 +router.post('/embed', async function (request, response) {
 +    try {
 +        const { server, items } = request.body;
++
 +        if (!server) {
 +            console.warn('KoboldCpp URL is not set');
 +            return response.sendStatus(400);
 +        }
++
 +        const headers = {};
 +        setAdditionalHeadersByType(headers, TEXTGEN_TYPES.KOBOLDCPP, server, request.user.directories);
++
 +        const embeddingsUrl = new URL(server);
 +        embeddingsUrl.pathname = '/api/extra/embeddings';
++
 +        const embeddingsResult = await fetch(embeddingsUrl, {
 +            method: 'POST',
 +            headers: {
 +                ...headers,
 +            },
 +            body: JSON.stringify({
 +                input: items,
 +            }),
 +        });
++
 +        /** @type {any} */
 +        const data = await embeddingsResult.json();
++
 +        if (!Array.isArray(data?.data)) {
 +            console.warn('KoboldCpp API response was not an array');
 +            return response.sendStatus(500);
 +        }
++
 +        const model = data.model || 'unknown';
 +        const embeddings = data.data.map(x => Array.isArray(x) ? x[0] : x).sort((a, b) => a.index - b.index).map(x => x.embedding);
 +        return response.json({ model, embeddings });
 +    } catch (error) {
 +        console.error('KoboldCpp embedding failed', error);
 +        response.status(500).send('Internal server error');
 +    }
 +});

src/endpoints/vectors.js+11 -0 Viewed

     'llamacpp',
     'vllm',
     'webllm',
 +    'koboldcpp',
 ];
 /**

             return getOllamaVector(text, sourceSettings.apiUrl, sourceSettings.model, sourceSettings.keep, directories);
         case 'webllm':
             return sourceSettings.embeddings[text];
 +        case 'koboldcpp':
 +            return sourceSettings.embeddings[text];
+    }
     throw new Error(`Unknown vector source ${source}`);

             case 'webllm':
                 results.push(...texts.map(x => sourceSettings.embeddings[x]));
                 break;
 +            case 'koboldcpp':
 +                results.push(...texts.map(x => sourceSettings.embeddings[x]));
 +                break;
             default:
                 throw new Error(`Unknown vector source ${source}`);
+        }

                 model: String(request.body.model),
                 embeddings: request.body.embeddings ?? {},
             };
 +        case 'koboldcpp':
 +            return {
 +                model: String(request.body.model),
 +                embeddings: request.body.embeddings ?? {},
 +            };
         default:
             return {};
+    }