From 7063fce2af86a7f50ca64275885ccf2091b94115 Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Mon, 6 May 2024 19:26:20 +0300
Subject: [PATCH] Selectable openrouter providers

---
 public/index.html                          | 10 +++
 public/scripts/openai.js                   |  7 ++
 public/scripts/preset-manager.js           |  1 +
 public/scripts/textgen-models.js           | 75 +++++++++++++++++++++-
 public/scripts/textgen-settings.js         |  6 ++
 src/constants.js                           |  1 +
 src/endpoints/backends/chat-completions.js |  7 ++
 src/endpoints/backends/text-completions.js |  8 +++
 8 files changed, 114 insertions(+), 1 deletion(-)

diff --git a/public/index.html b/public/index.html
index b93346a79..1c5d7049e 100644
--- a/public/index.html
+++ b/public/index.html
@@ -1989,6 +1989,11 @@
                                         </option>
                                     </select>
                                 </div>
+                                <div>
+                                    <h4 data-i18n="Model Providers">Model Providers</h4>
+                                    <select id="openrouter_providers_text" class="openrouter_providers" multiple>
+                                    </select>
+                                </div>
                             </div>
                             <div data-tg-type="infermaticai" class="flex-container flexFlowColumn">
                                 <h4 data-i18n="InfermaticAI API Key">InfermaticAI API Key</h4>
@@ -2503,6 +2508,11 @@
                                     <option data-i18n="Connect to the API">-- Connect to the API --</option>
                                 </select>
                             </div>
+                            <div>
+                                <h4 data-i18n="Model Providers">Model Providers</h4>
+                                <select id="openrouter_providers_chat" class="openrouter_providers" multiple>
+                                </select>
+                            </div>
                             <div class="marginTopBot5">
                                 <div class="inline-drawer wide100p">
                                     <div class="inline-drawer-toggle inline-drawer-header">
diff --git a/public/scripts/openai.js b/public/scripts/openai.js
index 094455bfc..a021ca0ee 100644
--- a/public/scripts/openai.js
+++ b/public/scripts/openai.js
@@ -254,6 +254,7 @@ const default_settings = {
     openrouter_force_instruct: false,
     openrouter_group_models: false,
     openrouter_sort_models: 'alphabetically',
+    openrouter_providers: [],
     jailbreak_system: false,
     reverse_proxy: '',
     chat_completion_source: chat_completion_sources.OPENAI,
@@ -328,6 +329,7 @@ const oai_settings = {
     openrouter_force_instruct: false,
     openrouter_group_models: false,
     openrouter_sort_models: 'alphabetically',
+    openrouter_providers: [],
     jailbreak_system: false,
     reverse_proxy: '',
     chat_completion_source: chat_completion_sources.OPENAI,
@@ -1770,6 +1772,7 @@ async function sendOpenAIRequest(type, messages, signal) {
         generate_data['repetition_penalty'] = Number(oai_settings.repetition_penalty_openai);
         generate_data['top_a'] = Number(oai_settings.top_a_openai);
         generate_data['use_fallback'] = oai_settings.openrouter_use_fallback;
+        generate_data['provider'] = oai_settings.openrouter_providers;
 
         if (isTextCompletion) {
             generate_data['stop'] = getStoppingStrings(isImpersonate, isContinue);
@@ -2793,6 +2796,7 @@ function loadOpenAISettings(data, settings) {
     $('#openrouter_use_fallback').prop('checked', oai_settings.openrouter_use_fallback);
     $('#openrouter_force_instruct').prop('checked', oai_settings.openrouter_force_instruct);
     $('#openrouter_group_models').prop('checked', oai_settings.openrouter_group_models);
+    $('#openrouter_providers_chat').val(oai_settings.openrouter_providers).trigger('change');
     $('#squash_system_messages').prop('checked', oai_settings.squash_system_messages);
     $('#continue_prefill').prop('checked', oai_settings.continue_prefill);
     if (settings.impersonation_prompt !== undefined) oai_settings.impersonation_prompt = settings.impersonation_prompt;
@@ -3016,6 +3020,7 @@ async function saveOpenAIPreset(name, settings, triggerUi = true) {
         openrouter_force_instruct: settings.openrouter_force_instruct,
         openrouter_group_models: settings.openrouter_group_models,
         openrouter_sort_models: settings.openrouter_sort_models,
+        openrouter_providers: settings.openrouter_providers,
         ai21_model: settings.ai21_model,
         mistralai_model: settings.mistralai_model,
         cohere_model: settings.cohere_model,
@@ -3412,6 +3417,7 @@ function onSettingsPresetChange() {
         openrouter_force_instruct: ['#openrouter_force_instruct', 'openrouter_force_instruct', true],
         openrouter_group_models: ['#openrouter_group_models', 'openrouter_group_models', false],
         openrouter_sort_models: ['#openrouter_sort_models', 'openrouter_sort_models', false],
+        openrouter_providers: ['#openrouter_providers_chat', 'openrouter_providers', false],
         ai21_model: ['#model_ai21_select', 'ai21_model', false],
         mistralai_model: ['#model_mistralai_select', 'mistralai_model', false],
         cohere_model: ['#model_cohere_select', 'cohere_model', false],
@@ -3496,6 +3502,7 @@ function onSettingsPresetChange() {
 
         $('#chat_completion_source').trigger('change');
         $('#openai_logit_bias_preset').trigger('change');
+        $('#openrouter_providers_chat').trigger('change');
 
         saveSettingsDebounced();
         eventSource.emit(event_types.OAI_PRESET_CHANGED_AFTER);
diff --git a/public/scripts/preset-manager.js b/public/scripts/preset-manager.js
index 889c254cb..e846fd1d0 100644
--- a/public/scripts/preset-manager.js
+++ b/public/scripts/preset-manager.js
@@ -319,6 +319,7 @@ class PresetManager {
             'dreamgen_model',
             'openrouter_model',
             'max_tokens_second',
+            'openrouter_providers',
         ];
         const settings = Object.assign({}, getSettingsByApiId(this.apiId));
 
diff --git a/public/scripts/textgen-models.js b/public/scripts/textgen-models.js
index 54252c45f..c4a6c61aa 100644
--- a/public/scripts/textgen-models.js
+++ b/public/scripts/textgen-models.js
@@ -1,7 +1,8 @@
 import { isMobile } from './RossAscends-mods.js';
-import { amount_gen, callPopup, eventSource, event_types, getRequestHeaders, max_context, setGenerationParamsFromPreset } from '../script.js';
+import { amount_gen, callPopup, eventSource, event_types, getRequestHeaders, max_context, saveSettingsDebounced, setGenerationParamsFromPreset } from '../script.js';
 import { textgenerationwebui_settings as textgen_settings, textgen_types } from './textgen-settings.js';
 import { tokenizers } from './tokenizers.js';
+import { oai_settings } from './openai.js';
 
 let mancerModels = [];
 let togetherModels = [];
@@ -11,6 +12,36 @@ let vllmModels = [];
 let aphroditeModels = [];
 export let openRouterModels = [];
 
+/**
+ * List of OpenRouter providers.
+ * @type {string[]}
+ */
+const OPENROUTER_PROVIDERS = [
+    'OpenAI',
+    'Anthropic',
+    'HuggingFace',
+    'Google',
+    'Mancer',
+    'Mancer 2',
+    'Together',
+    'DeepInfra',
+    'Azure',
+    'Modal',
+    'AnyScale',
+    'Replicate',
+    'Perplexity',
+    'Recursal',
+    'Fireworks',
+    'Mistral',
+    'Groq',
+    'Cohere',
+    'Lepton',
+    'OctoAI',
+    'Novita',
+    'Lynn',
+    'Lynn 2',
+];
+
 export async function loadOllamaModels(data) {
     if (!Array.isArray(data)) {
         console.error('Invalid Ollama models data', data);
@@ -476,6 +507,33 @@ jQuery(function () {
     $('#vllm_model').on('change', onVllmModelSelect);
     $('#aphrodite_model').on('change', onAphroditeModelSelect);
 
+    const providersSelect = $('.openrouter_providers');
+    for (const provider of OPENROUTER_PROVIDERS) {
+        providersSelect.append($('<option>', {
+            value: provider,
+            text: provider,
+        }));
+    }
+
+    providersSelect.on('change', function () {
+        const selectedProviders = $(this).val();
+
+        // Not a multiple select?
+        if (!Array.isArray(selectedProviders)) {
+            return;
+        }
+
+        if ($(this).is('#openrouter_providers_text')) {
+            textgen_settings.openrouter_providers = selectedProviders;
+        }
+
+        if ($(this).is('#openrouter_providers_chat')) {
+            oai_settings.openrouter_providers = selectedProviders;
+        }
+
+        saveSettingsDebounced();
+    });
+
     if (!isMobile()) {
         $('#mancer_model').select2({
             placeholder: 'Select a model',
@@ -532,5 +590,20 @@ jQuery(function () {
             width: '100%',
             templateResult: getAphroditeModelTemplate,
         });
+        providersSelect.select2({
+            sorter: data => data.sort((a, b) => a.text.localeCompare(b.text)),
+            placeholder: 'Select providers. No selection = all providers.',
+            searchInputPlaceholder: 'Search providers...',
+            searchInputCssClass: 'text_pole',
+            width: '100%',
+        });
+        providersSelect.on('select2:select', function (/** @type {any} */ evt) {
+            const element = evt.params.data.element;
+            const $element = $(element);
+
+            $element.detach();
+            $(this).append($element);
+            $(this).trigger('change');
+        });
     }
 });
diff --git a/public/scripts/textgen-settings.js b/public/scripts/textgen-settings.js
index a35763905..d63cadd9d 100644
--- a/public/scripts/textgen-settings.js
+++ b/public/scripts/textgen-settings.js
@@ -148,6 +148,7 @@ const settings = {
     infermaticai_model: '',
     ollama_model: '',
     openrouter_model: 'openrouter/auto',
+    openrouter_providers: [],
     vllm_model: '',
     aphrodite_model: '',
     dreamgen_model: 'opus-v1-xl/text',
@@ -454,6 +455,7 @@ function loadTextGenSettings(data, loadedSettings) {
     }
 
     $('#textgen_type').val(settings.type);
+    $('#openrouter_providers_text').val(settings.openrouter_providers).trigger('change');
     showTypeSpecificControls(settings.type);
     BIAS_CACHE.delete(BIAS_KEY);
     displayLogitBias(settings.logit_bias, BIAS_KEY);
@@ -1066,6 +1068,10 @@ export function getTextGenGenerationData(finalPrompt, maxTokens, isImpersonate,
         //'prompt_logprobs': settings.prompt_log_probs_aphrodite,
     };
 
+    if (settings.type === OPENROUTER) {
+        params.provider = settings.openrouter_providers;
+    }
+
     if (settings.type === KOBOLDCPP) {
         params.grammar = settings.grammar_string;
     }
diff --git a/src/constants.js b/src/constants.js
index c1c66594a..7ca267f74 100644
--- a/src/constants.js
+++ b/src/constants.js
@@ -298,6 +298,7 @@ const OPENROUTER_KEYS = [
     'stream',
     'prompt',
     'stop',
+    'provider',
 ];
 
 // https://github.com/vllm-project/vllm/blob/0f8a91401c89ac0a8018def3756829611b57727f/vllm/entrypoints/openai/protocol.py#L220
diff --git a/src/endpoints/backends/chat-completions.js b/src/endpoints/backends/chat-completions.js
index 368325ac9..5728bcf81 100644
--- a/src/endpoints/backends/chat-completions.js
+++ b/src/endpoints/backends/chat-completions.js
@@ -845,6 +845,13 @@ router.post('/generate', jsonParser, function (request, response) {
             bodyParams['repetition_penalty'] = request.body.repetition_penalty;
         }
 
+        if (Array.isArray(request.body.provider) && request.body.provider.length > 0) {
+            bodyParams['provider'] = {
+                allow_fallbacks: true,
+                order: request.body.provider ?? [],
+            };
+        }
+
         if (request.body.use_fallback) {
             bodyParams['route'] = 'fallback';
         }
diff --git a/src/endpoints/backends/text-completions.js b/src/endpoints/backends/text-completions.js
index 3cb15a748..49324e443 100644
--- a/src/endpoints/backends/text-completions.js
+++ b/src/endpoints/backends/text-completions.js
@@ -289,6 +289,14 @@ router.post('/generate', jsonParser, async function (request, response) {
         }
 
         if (request.body.api_type === TEXTGEN_TYPES.OPENROUTER) {
+            if (Array.isArray(request.body.provider) && request.body.provider.length > 0) {
+                request.body.provider = {
+                    allow_fallbacks: true,
+                    order: request.body.provider,
+                };
+            } else {
+                delete request.body.provider;
+            }
             request.body = _.pickBy(request.body, (_, key) => OPENROUTER_KEYS.includes(key));
             args.body = JSON.stringify(request.body);
         }