From 7ec9996c1f55bf8fc92d05d0b0fa6e5b20b2d7f2 Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Wed, 13 Mar 2024 02:55:09 +0200
Subject: [PATCH] Add prompt cost for OpenRouter under text completion

---
 public/index.html                |  3 +++
 public/scripts/textgen-models.js | 31 ++++++++++++++++++++++++++++++-
 2 files changed, 33 insertions(+), 1 deletion(-)

diff --git a/public/index.html b/public/index.html
index 68d97862b..ae73b3b8d 100644
--- a/public/index.html
+++ b/public/index.html
@@ -266,6 +266,9 @@
                                         </label>
                                     </div>
                                 </div>
+                                <div data-tg-type="openrouter">
+                                    Max prompt cost: <span id="openrouter_max_prompt_cost_tg">Unknown</span>
+                                </div>
                             </div>
                             <hr>
                         </div>
diff --git a/public/scripts/textgen-models.js b/public/scripts/textgen-models.js
index 51e98dee8..230c415b7 100644
--- a/public/scripts/textgen-models.js
+++ b/public/scripts/textgen-models.js
@@ -1,5 +1,5 @@
 import { isMobile } from './RossAscends-mods.js';
-import { callPopup, getRequestHeaders, setGenerationParamsFromPreset } from '../script.js';
+import { amount_gen, callPopup, eventSource, event_types, getRequestHeaders, max_context, setGenerationParamsFromPreset } from '../script.js';
 import { textgenerationwebui_settings as textgen_settings, textgen_types } from './textgen-settings.js';
 import { tokenizers } from './tokenizers.js';
 
@@ -151,6 +151,9 @@ export async function loadOpenRouterModels(data) {
         option.selected = model.id === textgen_settings.openrouter_model;
         $('#openrouter_model').append(option);
     }
+
+    // Calculate the cost of the selected model + update on settings change
+    calculateOpenRouterCost();
 }
 
 export async function loadAphroditeModels(data) {
@@ -362,6 +365,32 @@ async function downloadOllamaModel() {
     }
 }
 
+function calculateOpenRouterCost() {
+    if (textgen_settings.type !== textgen_types.OPENROUTER) {
+        return;
+    }
+
+    let cost = 'Unknown';
+    const model = openRouterModels.find(x => x.id === textgen_settings.openrouter_model);
+
+    if (model?.pricing) {
+        const completionCost = Number(model.pricing.completion);
+        const promptCost = Number(model.pricing.prompt);
+        const completionTokens = amount_gen;
+        const promptTokens = (max_context - completionTokens);
+        const totalCost = (completionCost * completionTokens) + (promptCost * promptTokens);
+        if (!isNaN(totalCost)) {
+            cost = '$' + totalCost.toFixed(3);
+        }
+    }
+
+    $('#openrouter_max_prompt_cost_tg').text(cost);
+
+    // Schedule an update when settings change
+    eventSource.removeListener(event_types.SETTINGS_UPDATED, calculateOpenRouterCost);
+    eventSource.once(event_types.SETTINGS_UPDATED, calculateOpenRouterCost);
+}
+
 export function getCurrentOpenRouterModelTokenizer() {
     const modelId = textgen_settings.openrouter_model;
     const model = openRouterModels.find(x => x.id === modelId);