From c53e496687521e5c5b5e6303d0dd2e7aa18ab5f1 Mon Sep 17 00:00:00 2001
From: Cohee <titan0599@gmail.com>
Date: Tue, 11 Jul 2023 15:46:41 +0300
Subject: [PATCH] Split OpenRouter and Window AI sources

---
 public/index.html                  |  22 +--
 public/scripts/RossAscends-mods.js |   4 +-
 public/scripts/openai.js           | 218 +++++++++++++++++++----------
 3 files changed, 161 insertions(+), 83 deletions(-)

diff --git a/public/index.html b/public/index.html
index 27bbd3108..b1f327044 100644
--- a/public/index.html
+++ b/public/index.html
@@ -564,7 +564,7 @@
                                         <input type="number" id="openai_max_tokens" name="openai_max_tokens" class="text_pole" min="50" max="1000">
                                     </div>
                                 </div>
-                                <div class="range-block" data-source="openai,claude,windowai">
+                                <div class="range-block" data-source="openai,claude,windowai,openrouter">
                                     <div class="range-block-title" data-i18n="Temperature">
                                         Temperature
                                     </div>
@@ -1340,7 +1340,7 @@
                             <option value="koboldhorde"><span data-i18n="KoboldAI Horde">KoboldAI Horde</span></option>
                             <option value="textgenerationwebui"><span data-i18n="Text Gen WebUI (ooba)">Text Gen WebUI (ooba)</span></option>
                             <option value="novel"><span data-i18n="NovelAI">NovelAI</span></option>
-                            <option value="openai"><span data-i18n="Chat Completion (OpenAI, Claude, Window/OpenRouter, Scale)">Chat Completion (OpenAI, Claude, Window/OpenRouter, Scale)</span></option>
+                            <option value="openai"><span data-i18n="Chat Completion (OpenAI, Claude, Window/OpenRouter, Scale)">Chat Completion (OpenAI, Claude, Window, OpenRouter, Scale)</span></option>
                             <option value="poe"><span data-i18n="Poe">Poe</span></option>
                         </select>
                     </div>
@@ -1497,7 +1497,8 @@
                         </h3>
                         <select id="chat_completion_source">
                             <option value="openai">OpenAI</option>
-                            <option value="windowai">Window AI / OpenRouter</option>
+                            <option value="windowai">Window AI</option>
+                            <option value="openrouter">OpenRouter</option>
                             <option value="claude">Claude</option>
                             <option value="scale">Scale</option>
                         </select>
@@ -1612,6 +1613,15 @@
                                     <option value="cohere/command-nightly">cohere/command-nightly</option>
                                 </select>
                             </div>
+                        </form>
+
+                        <form id="openrouter_form" data-source="openrouter" action="javascript:void(null);" method="post" enctype="multipart/form-data">
+                            <div>
+                                <h4 data-i18n="OpenRouter Model">OpenRouter Model</h4>
+                                <select id="model_openrouter_select">
+                                    <option data-i18n="Connect to the API">-- Connect to the API --</option>
+                                </select>
+                            </div>
                             <h4 data-i18n="OpenRouter API Key">OpenRouter API Key</h4>
                             <div>
                                 <small>
@@ -1625,10 +1635,6 @@
                             <div data-for="api_key_openrouter" class="neutral_warning">
                                 For privacy reasons, your API key will be hidden after you reload the page.
                             </div>
-                            <label class="checkbox_label" for="use_openrouter">
-                                <input id="use_openrouter" type="checkbox" />
-                                <span data-i18n="Use OpenRouter">Use OpenRouter directly</span>
-                            </label>
                         </form>
 
                         <form id="scale_form" data-source="scale" action="javascript:void(null);" method="post" enctype="multipart/form-data">
@@ -1648,7 +1654,7 @@
 
                         <div class="flex-container flex">
                             <input id="api_button_openai" class="menu_button" type="submit" value="Connect">
-                            <input data-source="windowai" id="openrouter_authorize" class="menu_button" type="button" value="Authorize" title="Get your OpenRouter API token using OAuth flow. You will be redirected to openrouter.ai">
+                            <input data-source="openrouter" id="openrouter_authorize" class="menu_button" type="button" value="Authorize" title="Get your OpenRouter API token using OAuth flow. You will be redirected to openrouter.ai">
                             <input id="test_api_button" class="menu_button" type="button" value="Test Message" title="Verifies your API connection by sending a short test message. Be aware that you'll be credited for it!">
                         </div>
                         <div id="api_loading_openai" class=" api-load-icon fa-solid fa-hourglass fa-spin"></div>
diff --git a/public/scripts/RossAscends-mods.js b/public/scripts/RossAscends-mods.js
index a892d6dd0..952c51613 100644
--- a/public/scripts/RossAscends-mods.js
+++ b/public/scripts/RossAscends-mods.js
@@ -406,8 +406,8 @@ function RA_autoconnect(PrevApi) {
                 if (   (secret_state[SECRET_KEYS.OPENAI] && oai_settings.chat_completion_source == chat_completion_sources.OPENAI)
                     || (secret_state[SECRET_KEYS.CLAUDE] && oai_settings.chat_completion_source == chat_completion_sources.CLAUDE)
                     || (secret_state[SECRET_KEYS.SCALE] && oai_settings.chat_completion_source == chat_completion_sources.SCALE)
-                    || (oai_settings.chat_completion_source == chat_completion_sources.WINDOWAI && !oai_settings.use_openrouter)
-                    || (secret_state[SECRET_KEYS.OPENROUTER] && oai_settings.chat_completion_source == chat_completion_sources.WINDOWAI && oai_settings.use_openrouter)
+                    || (oai_settings.chat_completion_source == chat_completion_sources.WINDOWAI)
+                    || (secret_state[SECRET_KEYS.OPENROUTER] && oai_settings.chat_completion_source == chat_completion_sources.OPENROUTER)
                     ) {
                     $("#api_button_openai").click();
                 }
diff --git a/public/scripts/openai.js b/public/scripts/openai.js
index 151d85cf6..038f50c92 100644
--- a/public/scripts/openai.js
+++ b/public/scripts/openai.js
@@ -96,6 +96,7 @@ const oai_max_temp = 2.0;
 const claude_max_temp = 1.0;
 
 let biasCache = undefined;
+let model_list = [];
 const tokenCache = {};
 
 export const chat_completion_sources = {
@@ -103,6 +104,7 @@ export const chat_completion_sources = {
     WINDOWAI: 'windowai',
     CLAUDE: 'claude',
     SCALE: 'scale',
+    OPENROUTER: 'openrouter',
 };
 
 const default_settings = {
@@ -131,12 +133,12 @@ const default_settings = {
     openai_model: 'gpt-3.5-turbo',
     claude_model: 'claude-instant-v1',
     windowai_model: '',
+    openrouter_model: 'openai/gpt-3.5-turbo',
     jailbreak_system: false,
     reverse_proxy: '',
     legacy_streaming: false,
     chat_completion_source: chat_completion_sources.OPENAI,
     max_context_unlocked: false,
-    use_openrouter: false,
     api_url_scale: '',
 };
 
@@ -166,12 +168,12 @@ const oai_settings = {
     openai_model: 'gpt-3.5-turbo',
     claude_model: 'claude-instant-v1',
     windowai_model: '',
+    openrouter_model: 'openai/gpt-3.5-turbo',
     jailbreak_system: false,
     reverse_proxy: '',
     legacy_streaming: false,
     chat_completion_source: chat_completion_sources.OPENAI,
     max_context_unlocked: false,
-    use_openrouter: false,
     api_url_scale: '',
 };
 
@@ -675,11 +677,35 @@ function getChatCompletionModel() {
             return oai_settings.windowai_model;
         case chat_completion_sources.SCALE:
             return '';
+        case chat_completion_sources.OPENROUTER:
+            return oai_settings.openrouter_model;
         default:
             throw new Error(`Unknown chat completion source: ${oai_settings.chat_completion_source}`);
     }
 }
 
+function saveModelList(data) {
+    model_list = data.map((model) => ({ id: model.id, context_length: model.context_length }));
+
+    if (oai_settings.chat_completion_source == chat_completion_sources.OPENROUTER) {
+        $('#model_openrouter_select').empty();
+        model_list.forEach((model) => {
+            const selected = model.id == oai_settings.openrouter_model;
+            $('#model_openrouter_select').append(
+                $('<option>', {
+                    value: model.id,
+                    text: model.id,
+                    selected: selected,
+                }));
+            if (selected) {
+                $('#model_openrouter_select').val(model.id).trigger('change');
+            }
+        });
+    }
+
+    // TODO Add ability to select OpenAI model from endpoint-provided list
+}
+
 async function sendOpenAIRequest(type, openai_msgs_tosend, signal) {
     // Provide default abort signal
     if (!signal) {
@@ -688,14 +714,14 @@ async function sendOpenAIRequest(type, openai_msgs_tosend, signal) {
 
     let logit_bias = {};
     const isClaude = oai_settings.chat_completion_source == chat_completion_sources.CLAUDE;
-    const isOpenRouter = oai_settings.use_openrouter && oai_settings.chat_completion_source == chat_completion_sources.WINDOWAI;
+    const isOpenRouter = oai_settings.chat_completion_source == chat_completion_sources.OPENROUTER;
     const isScale = oai_settings.chat_completion_source == chat_completion_sources.SCALE;
     const isTextCompletion = oai_settings.chat_completion_source == chat_completion_sources.OPENAI && (oai_settings.openai_model.startsWith('text-') || oai_settings.openai_model.startsWith('code-'));
     const stream = type !== 'quiet' && oai_settings.stream_openai && !isScale;
 
     // If we're using the window.ai extension, use that instead
     // Doesn't support logit bias yet
-    if (oai_settings.chat_completion_source == chat_completion_sources.WINDOWAI && !oai_settings.use_openrouter) {
+    if (oai_settings.chat_completion_source == chat_completion_sources.WINDOWAI) {
         return sendWindowAIRequest(openai_msgs_tosend, signal, stream);
     }
 
@@ -1030,8 +1056,8 @@ function loadOpenAISettings(data, settings) {
     oai_settings.wi_format = settings.wi_format ?? default_settings.wi_format;
     oai_settings.claude_model = settings.claude_model ?? default_settings.claude_model;
     oai_settings.windowai_model = settings.windowai_model ?? default_settings.windowai_model;
+    oai_settings.openrouter_model = settings.openrouter_model ?? default_settings.openrouter_model;
     oai_settings.chat_completion_source = settings.chat_completion_source ?? default_settings.chat_completion_source;
-    oai_settings.use_openrouter = settings.use_openrouter ?? default_settings.use_openrouter;
     oai_settings.api_url_scale = settings.api_url_scale ?? default_settings.api_url_scale;
 
     if (settings.nsfw_toggle !== undefined) oai_settings.nsfw_toggle = !!settings.nsfw_toggle;
@@ -1053,6 +1079,7 @@ function loadOpenAISettings(data, settings) {
     $(`#model_windowai_select option[value="${oai_settings.windowai_model}"`).attr('selected', true);
     $('#openai_max_context').val(oai_settings.openai_max_context);
     $('#openai_max_context_counter').text(`${oai_settings.openai_max_context}`);
+    $('#model_openrouter_select').val(oai_settings.openrouter_model);
 
     $('#openai_max_tokens').val(oai_settings.openai_max_tokens);
 
@@ -1110,12 +1137,11 @@ function loadOpenAISettings(data, settings) {
 
     $('#chat_completion_source').val(oai_settings.chat_completion_source).trigger('change');
     $('#oai_max_context_unlocked').prop('checked', oai_settings.max_context_unlocked);
-    $('#use_openrouter').prop('checked', oai_settings.use_openrouter);
 }
 
 async function getStatusOpen() {
     if (is_get_status_openai) {
-        if (oai_settings.chat_completion_source == chat_completion_sources.WINDOWAI && !oai_settings.use_openrouter) {
+        if (oai_settings.chat_completion_source == chat_completion_sources.WINDOWAI) {
             let status;
 
             if ('ai' in window) {
@@ -1138,7 +1164,7 @@ async function getStatusOpen() {
 
         let data = {
             reverse_proxy: oai_settings.reverse_proxy,
-            use_openrouter: oai_settings.use_openrouter && oai_settings.chat_completion_source == chat_completion_sources.WINDOWAI,
+            use_openrouter: oai_settings.chat_completion_source == chat_completion_sources.OPENROUTER,
         };
 
         return jQuery.ajax({
@@ -1156,6 +1182,9 @@ async function getStatusOpen() {
             success: function (data) {
                 if (!('error' in data))
                     setOnlineStatus('Valid');
+                if ('data' in data && Array.isArray(data.data)) {
+                    saveModelList(data.data);
+                }
                 resultCheckStatusOpen();
             },
             error: function (jqXHR, exception) {
@@ -1214,7 +1243,7 @@ async function saveOpenAIPreset(name, settings) {
         openai_model: settings.openai_model,
         claude_model: settings.claude_model,
         windowai_model: settings.windowai_model,
-        use_openrouter: settings.use_openrouter,
+        openrouter_model: settings.openrouter_model,
         temperature: settings.temp_openai,
         frequency_penalty: settings.freq_pen_openai,
         presence_penalty: settings.pres_pen_openai,
@@ -1571,6 +1600,7 @@ function onSettingsPresetChange() {
         openai_model: ['#model_openai_select', 'openai_model', false],
         claude_model: ['#model_claude_select', 'claude_model', false],
         windowai_model: ['#model_windowai_select', 'windowai_model', false],
+        openrouter_model: ['#model_openrouter_select', 'openrouter_model', false],
         openai_max_context: ['#openai_max_context', 'openai_max_context', false],
         openai_max_tokens: ['#openai_max_tokens', 'openai_max_tokens', false],
         nsfw_toggle: ['#nsfw_toggle', 'nsfw_toggle', true],
@@ -1589,7 +1619,6 @@ function onSettingsPresetChange() {
         nsfw_avoidance_prompt: ['#nsfw_avoidance_prompt_textarea', 'nsfw_avoidance_prompt', false],
         wi_format: ['#wi_format_textarea', 'wi_format', false],
         stream_openai: ['#stream_toggle', 'stream_openai', true],
-        use_openrouter: ['#use_openrouter', 'use_openrouter', true],
         api_url_scale: ['#api_url_scale', 'api_url_scale', false],
     };
 
@@ -1609,6 +1638,66 @@ function onSettingsPresetChange() {
     saveSettingsDebounced();
 }
 
+function getMaxContextOpenAI(value) {
+    if (oai_settings.max_context_unlocked) {
+        return unlocked_max;
+    }
+    else if (['gpt-4', 'gpt-4-0314', 'gpt-4-0613'].includes(value)) {
+        return max_8k;
+    }
+    else if (['gpt-4-32k', 'gpt-4-32k-0314', 'gpt-4-32k-0613'].includes(value)) {
+        return max_32k;
+    }
+    else if (['gpt-3.5-turbo-16k', 'gpt-3.5-turbo-16k-0613'].includes(value)) {
+        return max_16k;
+    }
+    else if (value == 'code-davinci-002') {
+        return max_8k;
+    }
+    else if (['text-curie-001', 'text-babbage-001', 'text-ada-001'].includes(value)) {
+        return max_2k;
+    }
+    else {
+        // default to gpt-3 (4095 tokens)
+        return max_4k;
+    }
+}
+
+
+function getMaxContextWindowAI(value) {
+    if (oai_settings.max_context_unlocked) {
+        return unlocked_max;
+    }
+    else if (value.endsWith('100k')) {
+        return claude_100k_max;
+    }
+    else if (value.includes('claude')) {
+        return claude_max;
+    }
+    else if (value.includes('gpt-3.5-turbo-16k')) {
+        return max_16k;
+    }
+    else if (value.includes('gpt-3.5')) {
+        return max_4k;
+    }
+    else if (value.includes('gpt-4-32k')) {
+        return max_32k;
+    }
+    else if (value.includes('gpt-4')) {
+        return max_8k;
+    }
+    else if (value.includes('palm-2')) {
+        return palm2_max;
+    }
+    else if (value.includes('GPT-NeoXT')) {
+        return max_2k;
+    }
+    else {
+        // default to gpt-3 (4095 tokens)
+        return max_4k;
+    }
+}
+
 async function onModelChange() {
     let value = $(this).val();
 
@@ -1627,6 +1716,16 @@ async function onModelChange() {
         oai_settings.openai_model = value;
     }
 
+    if ($(this).is('#model_openrouter_select')) {
+        if (!value) {
+            console.debug('Null OR model selected. Ignoring.');
+            return;
+        }
+
+        console.log('OpenRouter model changed to', value);
+        oai_settings.openrouter_model = value;
+    }
+
     if (oai_settings.chat_completion_source == chat_completion_sources.SCALE) {
         if (oai_settings.max_context_unlocked) {
             $('#openai_max_context').attr('max', unlocked_max);
@@ -1637,6 +1736,30 @@ async function onModelChange() {
         $('#openai_max_context').val(oai_settings.openai_max_context).trigger('input');
     }
 
+    if (oai_settings.chat_completion_source == chat_completion_sources.OPENROUTER) {
+        if (oai_settings.max_context_unlocked) {
+            $('#openai_max_context').attr('max', unlocked_max);
+        } else {
+            const model = model_list.find(m => m.id == oai_settings.openrouter_model);
+            if (model?.context_length) {
+                $('#openai_max_context').attr('max', model.context_length - 1); // waiting for openrouter to fix this
+            } else {
+                $('#openai_max_context').attr('max', max_4k); // placeholder
+            }
+        }
+        oai_settings.openai_max_context = Math.min(Number($('#openai_max_context').attr('max')), oai_settings.openai_max_context);
+        $('#openai_max_context').val(oai_settings.openai_max_context).trigger('input');
+
+        if (value && (value.includes('claude') || value.includes('palm-2'))) {
+            oai_settings.temp_openai = Math.min(claude_max_temp, oai_settings.temp_openai);
+            $('#temp_openai').attr('max', claude_max_temp).val(oai_settings.temp_openai).trigger('input');
+        }
+        else {
+            oai_settings.temp_openai = Math.min(oai_max_temp, oai_settings.temp_openai);
+            $('#temp_openai').attr('max', oai_max_temp).val(oai_settings.temp_openai).trigger('input');
+        }
+    }
+
     if (oai_settings.chat_completion_source == chat_completion_sources.CLAUDE) {
         if (oai_settings.max_context_unlocked) {
             $('#openai_max_context').attr('max', unlocked_max);
@@ -1662,38 +1785,7 @@ async function onModelChange() {
             value = (await window.ai.getCurrentModel()) || '';
         }
 
-        if (oai_settings.max_context_unlocked) {
-            $('#openai_max_context').attr('max', unlocked_max);
-        }
-        else if (value.endsWith('100k')) {
-            $('#openai_max_context').attr('max', claude_100k_max);
-        }
-        else if (value.includes('claude')) {
-            $('#openai_max_context').attr('max', claude_max);
-        }
-        else if (value.includes('gpt-3.5-turbo-16k')) {
-            $('#openai_max_context').attr('max', max_16k);
-        }
-        else if (value.includes('gpt-3.5')) {
-            $('#openai_max_context').attr('max', max_4k);
-        }
-        else if (value.includes('gpt-4-32k')) {
-            $('#openai_max_context').attr('max', max_32k);
-        }
-        else if (value.includes('gpt-4')) {
-            $('#openai_max_context').attr('max', max_8k);
-        }
-        else if (value.includes('palm-2')) {
-            $('#openai_max_context').attr('max', palm2_max);
-        }
-        else if (value.includes('GPT-NeoXT')) {
-            $('#openai_max_context').attr('max', max_2k);
-        }
-        else {
-            // default to gpt-3 (4095 tokens)
-            $('#openai_max_context').attr('max', max_4k);
-        }
-
+        $('#openai_max_context').attr('max', getMaxContextWindowAI(value));
         oai_settings.openai_max_context = Math.min(Number($('#openai_max_context').attr('max')), oai_settings.openai_max_context);
         $('#openai_max_context').val(oai_settings.openai_max_context).trigger('input');
 
@@ -1708,28 +1800,7 @@ async function onModelChange() {
     }
 
     if (oai_settings.chat_completion_source == chat_completion_sources.OPENAI) {
-        if (oai_settings.max_context_unlocked) {
-            $('#openai_max_context').attr('max', unlocked_max);
-        }
-        else if (['gpt-4', 'gpt-4-0314', 'gpt-4-0613'].includes(value)) {
-            $('#openai_max_context').attr('max', max_8k);
-        }
-        else if (['gpt-4-32k', 'gpt-4-32k-0314', 'gpt-4-32k-0613'].includes(value)) {
-            $('#openai_max_context').attr('max', max_32k);
-        }
-        else if (['gpt-3.5-turbo-16k', 'gpt-3.5-turbo-16k-0613'].includes(value)) {
-            $('#openai_max_context').attr('max', max_16k);
-        }
-        else if (value == 'code-davinci-002') {
-            $('#openai_max_context').attr('max', max_8k);
-        }
-        else if (['text-curie-001', 'text-babbage-001', 'text-ada-001'].includes(value)) {
-            $('#openai_max_context').attr('max', max_2k);
-        }
-        else {
-            $('#openai_max_context').attr('max', max_4k);
-        }
-
+        $('#openai_max_context').attr('max', getMaxContextOpenAI(value));
         oai_settings.openai_max_context = Math.min(oai_settings.openai_max_context, Number($('#openai_max_context').attr('max')));
         $('#openai_max_context').val(oai_settings.openai_max_context).trigger('input');
 
@@ -1769,18 +1840,21 @@ async function onConnectButtonClick(e) {
     if (oai_settings.chat_completion_source == chat_completion_sources.WINDOWAI) {
         is_get_status_openai = true;
         is_api_button_press_openai = true;
+
+        return await getStatusOpen();
+    }
+
+    if (oai_settings.chat_completion_source == chat_completion_sources.OPENROUTER) {
         const api_key_openrouter = $('#api_key_openrouter').val().trim();
 
         if (api_key_openrouter.length) {
             await writeSecret(SECRET_KEYS.OPENROUTER, api_key_openrouter);
         }
 
-        if (oai_settings.use_openrouter && !secret_state[SECRET_KEYS.OPENROUTER]) {
+        if (!secret_state[SECRET_KEYS.OPENROUTER]) {
             console.log('No secret key saved for OpenRouter');
             return;
         }
-
-        return await getStatusOpen();
     }
 
     if (oai_settings.chat_completion_source == chat_completion_sources.SCALE) {
@@ -1848,6 +1922,9 @@ function toggleChatCompletionForms() {
     else if (oai_settings.chat_completion_source == chat_completion_sources.SCALE) {
         $('#model_scale_select').trigger('change');
     }
+    else if (oai_settings.chat_completion_source == chat_completion_sources.OPENROUTER) {
+        $('#model_openrouter_select').trigger('change');
+    }
 
     $('[data-source]').each(function () {
         const validSources = $(this).data('source').split(',');
@@ -2074,12 +2151,6 @@ $(document).ready(function () {
         saveSettingsDebounced();
     });
 
-    $('#use_openrouter').on('input', function () {
-        oai_settings.use_openrouter = !!$(this).prop('checked');
-        reconnectOpenAi();
-        saveSettingsDebounced();
-    });
-
     $('#api_url_scale').on('input', function () {
         oai_settings.api_url_scale = $(this).val();
         saveSettingsDebounced();
@@ -2091,6 +2162,7 @@ $(document).ready(function () {
     $("#model_claude_select").on("change", onModelChange);
     $("#model_windowai_select").on("change", onModelChange);
     $("#model_scale_select").on("change", onModelChange);
+    $("#model_openrouter_select").on("change", onModelChange);
     $("#settings_perset_openai").on("change", onSettingsPresetChange);
     $("#new_oai_preset").on("click", onNewPresetClick);
     $("#delete_oai_preset").on("click", onDeletePresetClick);