From 57b81be9ce28f92cba7e50dc238204f1e48c889d Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Thu, 22 May 2025 23:03:04 +0300
Subject: [PATCH] Caption - allow custom endpoint for xAI

---
 public/scripts/extensions/caption/index.js    |  2 +-
 .../scripts/extensions/caption/settings.html  |  2 +-
 public/scripts/extensions/shared.js           |  4 ++--
 src/endpoints/openai.js                       | 20 +++++++++----------
 4 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/public/scripts/extensions/caption/index.js b/public/scripts/extensions/caption/index.js
index f29455af6..8e499452e 100644
--- a/public/scripts/extensions/caption/index.js
+++ b/public/scripts/extensions/caption/index.js
@@ -416,6 +416,7 @@ jQuery(async function () {
                         'google': SECRET_KEYS.MAKERSUITE,
                         'vertexai': SECRET_KEYS.VERTEXAI,
                         'anthropic': SECRET_KEYS.CLAUDE,
+                        'xai': SECRET_KEYS.XAI,
                     };
 
                     if (reverseProxyApis[api]) {
@@ -429,7 +430,6 @@ jQuery(async function () {
                         'zerooneai': SECRET_KEYS.ZEROONEAI,
                         'groq': SECRET_KEYS.GROQ,
                         'cohere': SECRET_KEYS.COHERE,
-                        'xai': SECRET_KEYS.XAI,
                     };
 
                     if (chatCompletionApis[api] && secret_state[chatCompletionApis[api]]) {
diff --git a/public/scripts/extensions/caption/settings.html b/public/scripts/extensions/caption/settings.html
index 79ddd5d81..9eb60ba5b 100644
--- a/public/scripts/extensions/caption/settings.html
+++ b/public/scripts/extensions/caption/settings.html
@@ -181,7 +181,7 @@
                 <div data-type="ollama">
                     The model must be downloaded first! Do it with the <code>ollama pull</code> command or <a href="#" id="caption_ollama_pull">click here</a>.
                 </div>
-                <label data-type="openai,anthropic,google,vertexai,mistral" class="checkbox_label flexBasis100p" for="caption_allow_reverse_proxy" title="Allow using reverse proxy if defined and valid.">
+                <label data-type="openai,anthropic,google,vertexai,mistral,xai" class="checkbox_label flexBasis100p" for="caption_allow_reverse_proxy" title="Allow using reverse proxy if defined and valid.">
                     <input id="caption_allow_reverse_proxy" type="checkbox" class="checkbox">
                     <span data-i18n="Allow reverse proxy">Allow reverse proxy</span>
                 </label>
diff --git a/public/scripts/extensions/shared.js b/public/scripts/extensions/shared.js
index e3ec2c29e..ad4e3589d 100644
--- a/public/scripts/extensions/shared.js
+++ b/public/scripts/extensions/shared.js
@@ -15,7 +15,7 @@ import { createThumbnail, isValidUrl } from '../utils.js';
  */
 export async function getMultimodalCaption(base64Img, prompt) {
     const useReverseProxy =
-        (['openai', 'anthropic', 'google', 'mistral', 'vertexai'].includes(extension_settings.caption.multimodal_api))
+        (['openai', 'anthropic', 'google', 'mistral', 'vertexai', 'xai'].includes(extension_settings.caption.multimodal_api))
         && extension_settings.caption.allow_reverse_proxy
         && oai_settings.reverse_proxy
         && isValidUrl(oai_settings.reverse_proxy);
@@ -157,7 +157,7 @@ function throwIfInvalidModel(useReverseProxy) {
         throw new Error('Cohere API key is not set.');
     }
 
-    if (extension_settings.caption.multimodal_api === 'xai' && !secret_state[SECRET_KEYS.XAI]) {
+    if (extension_settings.caption.multimodal_api === 'xai' && !secret_state[SECRET_KEYS.XAI] && !useReverseProxy) {
         throw new Error('xAI API key is not set.');
     }
 
diff --git a/src/endpoints/openai.js b/src/endpoints/openai.js
index 4e9d4e898..5af8e9976 100644
--- a/src/endpoints/openai.js
+++ b/src/endpoints/openai.js
@@ -22,8 +22,12 @@ router.post('/caption-image', async (request, response) => {
             key = readSecret(request.user.directories, SECRET_KEYS.OPENAI);
         }
 
-        if (request.body.api === 'openrouter' && !request.body.reverse_proxy) {
-            key = readSecret(request.user.directories, SECRET_KEYS.OPENROUTER);
+        if (request.body.api === 'xai' && !request.body.reverse_proxy) {
+            key = readSecret(request.user.directories, SECRET_KEYS.XAI);
+        }
+
+        if (request.body.api === 'mistral' && !request.body.reverse_proxy) {
+            key = readSecret(request.user.directories, SECRET_KEYS.MISTRALAI);
         }
 
         if (request.body.reverse_proxy && request.body.proxy_password) {
@@ -36,6 +40,10 @@ router.post('/caption-image', async (request, response) => {
             mergeObjectWithYaml(headers, request.body.custom_include_headers);
         }
 
+        if (request.body.api === 'openrouter') {
+            key = readSecret(request.user.directories, SECRET_KEYS.OPENROUTER);
+        }
+
         if (request.body.api === 'ooba') {
             key = readSecret(request.user.directories, SECRET_KEYS.OOBA);
             bodyParams.temperature = 0.1;
@@ -57,10 +65,6 @@ router.post('/caption-image', async (request, response) => {
             key = readSecret(request.user.directories, SECRET_KEYS.ZEROONEAI);
         }
 
-        if (request.body.api === 'mistral') {
-            key = readSecret(request.user.directories, SECRET_KEYS.MISTRALAI);
-        }
-
         if (request.body.api === 'groq') {
             key = readSecret(request.user.directories, SECRET_KEYS.GROQ);
         }
@@ -69,10 +73,6 @@ router.post('/caption-image', async (request, response) => {
             key = readSecret(request.user.directories, SECRET_KEYS.COHERE);
         }
 
-        if (request.body.api === 'xai') {
-            key = readSecret(request.user.directories, SECRET_KEYS.XAI);
-        }
-
         const noKeyTypes = ['custom', 'ooba', 'koboldcpp', 'vllm', 'llamacpp', 'pollinations'];
         if (!key && !request.body.reverse_proxy && !noKeyTypes.includes(request.body.api)) {
             console.warn('No key found for API', request.body.api);