feat: [Electron Hub] Support Vector Storage, Better searching for image engine (#4540) * feat: [Electron Hub] Add Vector Storage, Better searching for image model list * feat: [Electron Hub] Add quality parameter for Image Engine * fixed ESLint * Update public/scripts/extensions/vectors/index.js Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> * small tweaks * Use default getOpenAIVector * Refactor and clean-up code * Move endpoint filtering logic to backend --------- Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Cohee <18619528+Cohee1207@users.noreply.github.com>

f0ceba43e9364137dbbafda4d4afbbb86a057055

Ngo Dinh Gia Bao <112873281+snowby666@users.noreply.github.com>

Signed
Parents: 4add4f0090dc41d194c75c52684e9df6468b698f
8 files changed, +181 -6Ignore whitespace
public/scripts/extensions/stable-diffusion/index.js+77 -2
@@ -1277,9 +1277,16 @@ async function validateComfyUrl() {
12771277}
12781278
12791279async function onModelChange() {
12801280 extension_settings.sd.modelconst selectedModel = $('#sd_model').find(':selected').val();
1281+ extension_settings.sd.model = selectedModel.val();
12811282 saveSettingsDebounced();
12821283
1284+ if (extension_settings.sd.model && extension_settings.sd.source === sources.electronhub) {
1285+ const cachedModel = selectedModel.data('model');
1286+ const models = cachedModel ? [cachedModel] : await loadElectronHubModels();
1287+ ensureElectronHubQualitySelect(models);
1288+ }
1289+
12831290 const cloudSources = [
12841291 sources.horde,
12851292 sources.novel,
@@ -1726,11 +1733,16 @@ async function loadModels() {
17261733 break;
17271734 }
17281735
1736+ if (extension_settings.sd.source === sources.electronhub) {
1737+ ensureElectronHubQualitySelect(models);
1738+ }
1739+
17291740 for (const model of models) {
17301741 const option = document.createElement('option');
17311742 option.innerText = model.text;
17321743 option.value = model.value;
17331744 option.selected = model.value === extension_settings.sd.model;
1745+ $(option).data('model', model);
17341746 $('#sd_model').append(option);
17351747 }
17361748
@@ -1740,6 +1752,49 @@ async function loadModels() {
17401752 }
17411753}
17421754
1755+/**
1756+ * Ensure the Electron Hub quality select is populated based on the selected model.
1757+ * @param {any[]} models Array of models
1758+ */
1759+function ensureElectronHubQualitySelect(models) {
1760+ try {
1761+ const modelId = String(extension_settings.sd.model || '');
1762+ if (!modelId) return;
1763+
1764+ const model = Array.isArray(models) ? models.find(m => String(m?.id) === modelId) : undefined;
1765+ const qualities = Array.isArray(model?.qualities) ? model.qualities : undefined;
1766+
1767+ const $qualityRow = $('#sd_electronhub_quality_row');
1768+ const $select = $('#sd_electronhub_quality');
1769+
1770+ $qualityRow.toggle(!!qualities && qualities.length > 0);
1771+ $select.empty();
1772+
1773+ if (!qualities || qualities.length === 0) {
1774+ extension_settings.sd.electronhub_quality = undefined;
1775+ saveSettingsDebounced();
1776+ return;
1777+ }
1778+
1779+ for (const q of qualities) {
1780+ const opt = document.createElement('option');
1781+ opt.value = String(q);
1782+ opt.textContent = String(q);
1783+ opt.selected = String(q) === String(extension_settings.sd.electronhub_quality || '');
1784+ $select.append(opt);
1785+ }
1786+
1787+ if (!$select.val()) {
1788+ const first = String(qualities[0]);
1789+ extension_settings.sd.electronhub_quality = first;
1790+ $select.val(first);
1791+ saveSettingsDebounced();
1792+ }
1793+ } catch (e) {
1794+ console.error(e);
1795+ }
1796+}
1797+
17431798async function loadStabilityModels() {
17441799 $('#sd_stability_key').toggleClass('success', !!secret_state[SECRET_KEYS.STABILITY]);
17451800
@@ -1824,8 +1879,23 @@ async function loadElectronHubModels() {
18241879 headers: getRequestHeaders(),
18251880 });
18261881
1882+ function getModelName(model) {
1883+ const name = String(model?.name || model?.id || '');
1884+ const premium = model?.premium_model ? ' | Premium' : '';
1885+ let price = 'Unknown';
1886+ if (model?.pricing?.type === 'per_image') {
1887+ const coeff = Number(model.pricing.coefficient);
1888+ if (!isNaN(coeff)) {
1889+ price = `$${coeff}/image`;
1890+ }
1891+ }
1892+ return `${name} | ${price}${premium}`;
1893+ }
1894+
18271895 if (result.ok) {
1828- return await result.json();
1896+ /** @type {any[]} */
1897+ const data = await result.json();
1898+ return Array.isArray(data) ? data.map(m => ({ ...m, text: getModelName(m) })) : [];
18291899 }
18301900
18311901 return [];
@@ -3665,6 +3735,7 @@ async function generateElectronHubImage(prompt, signal) {
36653735 model: extension_settings.sd.model,
36663736 prompt: prompt,
36673737 size: size,
3738+ quality: String(extension_settings.sd.electronhub_quality || '').trim() || undefined,
36683739 }),
36693740 });
36703741
@@ -4816,6 +4887,10 @@ jQuery(async () => {
48164887 extension_settings.sd.google_enhance = $(this).prop('checked');
48174888 saveSettingsDebounced();
48184889 });
4890+ $('#sd_electronhub_quality').on('change', function () {
4891+ extension_settings.sd.electronhub_quality = String($(this).val());
4892+ saveSettingsDebounced();
4893+ });
48194894
48204895 if (!CSS.supports('field-sizing', 'content')) {
48214896 $('.sd_settings .inline-drawer-toggle').on('click', function () {
public/scripts/extensions/stable-diffusion/settings.html+6 -0
@@ -96,6 +96,12 @@
9696 </div>
9797 <div data-sd-source="electronhub">
9898 <i data-i18n="Hint: Save an API key in the Electron Hub (Chat Completion) API settings to use it here.">Hint: Save an API key in the Electron Hub (Chat Completion) API settings to use it here.</i>
99+ <div class="flex-container" id="sd_electronhub_quality_row">
100+ <div class="flex1">
101+ <label for="sd_electronhub_quality" data-i18n="Image Quality">Image Quality</label>
102+ <select id="sd_electronhub_quality"></select>
103+ </div>
104+ </div>
99105 </div>
100106 <div data-sd-source="nanogpt">
101107 <i data-i18n="Hint: Save an API key in the NanoGPT (Chat Completion) API settings to use it here.">Hint: Save an API key in the NanoGPT (Chat Completion) API settings to use it here.</i>
public/scripts/extensions/vectors/index.js+57 -2
@@ -61,6 +61,7 @@ const settings = {
6161 include_wi: false,
6262 togetherai_model: 'togethercomputer/m2-bert-80M-32k-retrieval',
6363 openai_model: 'text-embedding-ada-002',
64+ electronhub_model: 'text-embedding-3-small',
6465 cohere_model: 'embed-english-v3.0',
6566 ollama_model: 'mxbai-embed-large',
6667 ollama_keep: false,
@@ -771,6 +772,9 @@ function getVectorsRequestBody(args = {}) {
771772 body.extrasUrl = extension_settings.apiUrl;
772773 body.extrasKey = extension_settings.apiKey;
773774 break;
775+ case 'electronhub':
776+ body.model = extension_settings.vectors.electronhub_model;
777+ break;
774778 case 'togetherai':
775779 body.model = extension_settings.vectors.togetherai_model;
776780 break;
@@ -889,6 +893,7 @@ async function insertVectorItems(collectionId, items) {
889893 */
890894function throwIfSourceInvalid() {
891895 if (settings.source === 'openai' && !secret_state[SECRET_KEYS.OPENAI] ||
896+ settings.source === 'electronhub' && !secret_state[SECRET_KEYS.ELECTRONHUB] ||
892897 settings.source === 'palm' && !secret_state[SECRET_KEYS.MAKERSUITE] ||
893898 settings.source === 'vertexai' && !secret_state[SECRET_KEYS.VERTEXAI] && !secret_state[SECRET_KEYS.VERTEXAI_SERVICE_ACCOUNT] ||
894899 settings.source === 'mistral' && !secret_state[SECRET_KEYS.MISTRALAI] ||
@@ -1101,6 +1106,7 @@ function toggleSettings() {
11011106 $('#vectors_world_info_settings').toggle(!!settings.enabled_world_info);
11021107 $('#together_vectorsModel').toggle(settings.source === 'togetherai');
11031108 $('#openai_vectorsModel').toggle(settings.source === 'openai');
1109+ $('#electronhub_vectorsModel').toggle(settings.source === 'electronhub');
11041110 $('#cohere_vectorsModel').toggle(settings.source === 'cohere');
11051111 $('#ollama_vectorsModel').toggle(settings.source === 'ollama');
11061112 $('#llamacpp_vectorsModel').toggle(settings.source === 'llamacpp');
@@ -1110,12 +1116,56 @@ function toggleSettings() {
11101116 $('#koboldcpp_vectorsModel').toggle(settings.source === 'koboldcpp');
11111117 $('#google_vectorsModel').toggle(settings.source === 'palm' || settings.source === 'vertexai');
11121118 $('#vector_altEndpointUrl').toggle(vectorApiRequiresUrl.includes(settings.source));
11131119 ifswitch (settings.source === 'webllm') {
1114- loadWebLlmModels();
1120+ case 'webllm':
1121+ loadWebLlmModels();
1122+ break;
1123+ case 'electronhub':
1124+ loadElectronHubModels();
1125+ break;
1126+ }
1127+}
1128+
1129+async function loadElectronHubModels() {
1130+ try {
1131+ const response = await fetch('/api/openai/electronhub/models', {
1132+ method: 'POST',
1133+ headers: getRequestHeaders(),
1134+ });
1135+ if (!response.ok) {
1136+ throw new Error(`HTTP ${response.status}`);
1137+ }
1138+ /** @type {Array<any>} */
1139+ const data = await response.json();
1140+ // filter by embeddings endpoint
1141+ const models = Array.isArray(data) ? data.filter(m => Array.isArray(m?.endpoints) && m.endpoints.includes('/v1/embeddings')) : [];
1142+ populateElectronHubModelSelect(models);
1143+ } catch (err) {
1144+ console.warn('Electron Hub models fetch failed', err);
1145+ populateElectronHubModelSelect([]);
11151146 }
11161147}
11171148
11181149/**
1150+ * Populates the Electron Hub model select element.
1151+ * @param {{ id: string, name: string }[]} models Electron Hub models
1152+ */
1153+function populateElectronHubModelSelect(models) {
1154+ const select = $('#vectors_electronhub_model');
1155+ select.empty();
1156+ for (const m of models) {
1157+ const option = document.createElement('option');
1158+ option.value = m.id;
1159+ option.text = m.name || m.id;
1160+ select.append(option);
1161+ }
1162+ if (!settings.electronhub_model && models.length) {
1163+ settings.electronhub_model = models[0].id;
1164+ }
1165+ $('#vectors_electronhub_model').val(settings.electronhub_model);
1166+}
1167+
1168+/**
11191169 * Executes a function with WebLLM error handling.
11201170 * @param {function(): Promise<T>} func Function to execute
11211171 * @returns {Promise<T>}
@@ -1513,6 +1563,11 @@ jQuery(async () => {
15131563 Object.assign(extension_settings.vectors, settings);
15141564 saveSettingsDebounced();
15151565 });
1566+ $('#vectors_electronhub_model').val(settings.electronhub_model).on('change', () => {
1567+ settings.electronhub_model = String($('#vectors_electronhub_model').val());
1568+ Object.assign(extension_settings.vectors, settings);
1569+ saveSettingsDebounced();
1570+ });
15161571 $('#vectors_cohere_model').val(settings.cohere_model).on('change', () => {
15171572 settings.cohere_model = String($('#vectors_cohere_model').val());
15181573 Object.assign(extension_settings.vectors, settings);
public/scripts/extensions/vectors/settings.html+10 -0
@@ -11,6 +11,7 @@
1111 </label>
1212 <select id="vectors_source" class="text_pole">
1313 <option value="cohere">Cohere</option>
14+ <option value="electronhub">Electron Hub</option>
1415 <option value="extras">Extras (deprecated)</option>
1516 <option value="palm">Google AI Studio</option>
1617 <option value="vertexai">Google Vertex AI</option>
@@ -26,6 +27,15 @@
2627 <option value="webllm" data-i18n="WebLLM Extension">WebLLM Extension</option>
2728 </select>
2829 </div>
30+ <div class="flex-container flexFlowColumn" id="electronhub_vectorsModel">
31+ <label for="vectors_electronhub_model" data-i18n="Vectorization Model">
32+ Vectorization Model
33+ </label>
34+ <select id="vectors_electronhub_model" class="text_pole"></select>
35+ <i data-i18n="Hint: Set your Electron Hub API key in API Connections.">
36+ Hint: Set your Electron Hub API key in API Connections.
37+ </i>
38+ </div>
2939 <div class="flex-container flexFlowColumn" id="vector_altEndpointUrl">
3040 <label class="checkbox_label" for="vector_altEndpointUrl_enabled" title="Enable secondary endpoint URL usage, instead of the main one.">
3141 <input id="vector_altEndpointUrl_enabled" type="checkbox" class="checkbox">
public/scripts/openai.js+1 -1
@@ -4820,7 +4820,7 @@ function getElectronHubMaxContext(model, isUnlocked) {
48204820 return modelInfo.tokens;
48214821 }
48224822 }
48234823 return max_8kmax_128k;
48244824}
48254825
48264826/**
src/endpoints/stable-diffusion.js+15 -1
@@ -983,7 +983,15 @@ electronhub.post('/models', async (request, response) => {
983983
984984 /** @type {any} */
985985 const data = await modelsResponse.json();
986- const models = data.data.filter(x => x.endpoints.includes('/v1/images/generations')).map(x => ({ value: x.id, text: x.name }));
986+
987+ if (!Array.isArray(data?.data)) {
988+ console.warn('Electron Hub returned invalid data.');
989+ return response.sendStatus(500);
990+ }
991+
992+ const models = data.data
993+ .filter(x => x && Array.isArray(x.endpoints) && x.endpoints.includes('/v1/images/generations'))
994+ .map(x => ({ ...x, value: x.id, text: x.name }));
987995 return response.send(models);
988996 } catch (error) {
989997 console.error(error);
@@ -1010,6 +1018,12 @@ electronhub.post('/generate', async (request, response) => {
10101018 bodyParams.size = request.body.size;
10111019 }
10121020
1021+ if (request.body.quality) {
1022+ bodyParams.quality = request.body.quality;
1023+ }
1024+
1025+ console.debug('Electron Hub request:', bodyParams);
1026+
10131027 const result = await fetch('https://api.electronhub.ai/v1/images/generations', {
10141028 method: 'POST',
10151029 headers: {
src/endpoints/vectors.js+10 -0
@@ -34,6 +34,7 @@ const SOURCES = [
3434 'webllm',
3535 'koboldcpp',
3636 'vertexai',
37+ 'electronhub',
3738];
3839
3940/**
@@ -53,6 +54,8 @@ async function getVector(source, sourceSettings, text, isQuery, directories) {
5354 case 'mistral':
5455 case 'openai':
5556 return getOpenAIVector(text, source, directories, sourceSettings.model);
57+ case 'electronhub':
58+ return getOpenAIVector(text, source, directories, sourceSettings.model);
5659 case 'transformers':
5760 return getTransformersVector(text);
5861 case 'extras':
@@ -102,6 +105,9 @@ async function getBatchVector(source, sourceSettings, texts, isQuery, directorie
102105 case 'openai':
103106 results.push(...await getOpenAIBatchVector(batch, source, directories, sourceSettings.model));
104107 break;
108+ case 'electronhub':
109+ results.push(...await getOpenAIBatchVector(batch, source, directories, sourceSettings.model));
110+ break;
105111 case 'transformers':
106112 results.push(...await getTransformersBatchVector(batch));
107113 break;
@@ -156,6 +162,10 @@ function getSourceSettings(source, request) {
156162 return {
157163 model: String(request.body.model),
158164 };
165+ case 'electronhub':
166+ return {
167+ model: String(request.body.model || 'text-embedding-3-small'),
168+ };
159169 case 'cohere':
160170 return {
161171 model: String(request.body.model),
src/vectors/openai-vectors.js+5 -0
@@ -17,6 +17,11 @@ const SOURCES = {
1717 url: 'api.openai.com',
1818 model: 'text-embedding-ada-002',
1919 },
20+ 'electronhub': {
21+ secretKey: SECRET_KEYS.ELECTRONHUB,
22+ url: 'api.electronhub.ai',
23+ model: 'text-embedding-3-small',
24+ },
2025};
2126
2227/**