From e9cf606c70e2cba2aade710acfdfc6b6ca28bb3b Mon Sep 17 00:00:00 2001
From: Cohee <18619528+Cohee1207@users.noreply.github.com>
Date: Thu, 6 Mar 2025 22:23:35 +0200
Subject: [PATCH] Add backend-provided websearch connectors for OpenRouter and
 Gemini

---
 public/index.html                          | 11 ++++++++
 public/scripts/openai.js                   | 12 +++++++++
 src/endpoints/backends/chat-completions.js | 29 ++++++++++++++++++++++
 3 files changed, 52 insertions(+)

diff --git a/public/index.html b/public/index.html
index aab2a0c11..6724c399f 100644
--- a/public/index.html
+++ b/public/index.html
@@ -1951,6 +1951,17 @@
                                             </span>
                                         </div>
                                     </div>
+                                    <div class="range-block" data-source="makersuite,openrouter">
+                                        <label for="openai_enable_web_search" class="checkbox_label flexWrap widthFreeExpand">
+                                            <input id="openai_enable_web_search" type="checkbox" />
+                                            <span data-i18n="Enable web search">Enable web search</span>
+                                        </label>
+                                        <div class="flexBasis100p toggle-description justifyLeft">
+                                            <span>
+                                                Use search capabilities provided by the backend.
+                                            </span>
+                                        </div>
+                                    </div>
                                     <div class="range-block" data-source="openai,cohere,mistralai,custom,claude,openrouter,groq,deepseek">
                                         <label for="openai_function_calling" class="checkbox_label flexWrap widthFreeExpand">
                                             <input id="openai_function_calling" type="checkbox" />
diff --git a/public/scripts/openai.js b/public/scripts/openai.js
index 2ea38dc7b..33dc24978 100644
--- a/public/scripts/openai.js
+++ b/public/scripts/openai.js
@@ -300,6 +300,7 @@ export const settingsToUpdate = {
     function_calling: ['#openai_function_calling', 'function_calling', true],
     show_thoughts: ['#openai_show_thoughts', 'show_thoughts', true],
     reasoning_effort: ['#openai_reasoning_effort', 'reasoning_effort', false],
+    enable_web_search: ['#openai_enable_web_search', 'enable_web_search', true],
     seed: ['#seed_openai', 'seed', false],
     n: ['#n_openai', 'n', false],
     bypass_status_check: ['#openai_bypass_status_check', 'bypass_status_check', true],
@@ -380,6 +381,7 @@ const default_settings = {
     custom_prompt_post_processing: custom_prompt_post_processing_types.NONE,
     show_thoughts: true,
     reasoning_effort: 'medium',
+    enable_web_search: false,
     seed: -1,
     n: 1,
 };
@@ -459,6 +461,7 @@ const oai_settings = {
     custom_prompt_post_processing: custom_prompt_post_processing_types.NONE,
     show_thoughts: true,
     reasoning_effort: 'medium',
+    enable_web_search: false,
     seed: -1,
     n: 1,
 };
@@ -2000,6 +2003,7 @@ async function sendOpenAIRequest(type, messages, signal) {
         'group_names': getGroupNames(),
         'include_reasoning': Boolean(oai_settings.show_thoughts),
         'reasoning_effort': String(oai_settings.reasoning_effort),
+        'enable_web_search': Boolean(oai_settings.enable_web_search),
     };
 
     if (!canMultiSwipe && ToolManager.canPerformToolCalls(type)) {
@@ -3222,6 +3226,7 @@ function loadOpenAISettings(data, settings) {
     oai_settings.bypass_status_check = settings.bypass_status_check ?? default_settings.bypass_status_check;
     oai_settings.show_thoughts = settings.show_thoughts ?? default_settings.show_thoughts;
     oai_settings.reasoning_effort = settings.reasoning_effort ?? default_settings.reasoning_effort;
+    oai_settings.enable_web_search = settings.enable_web_search ?? default_settings.enable_web_search;
     oai_settings.seed = settings.seed ?? default_settings.seed;
     oai_settings.n = settings.n ?? default_settings.n;
 
@@ -3349,6 +3354,7 @@ function loadOpenAISettings(data, settings) {
     $('#seed_openai').val(oai_settings.seed);
     $('#n_openai').val(oai_settings.n);
     $('#openai_show_thoughts').prop('checked', oai_settings.show_thoughts);
+    $('#openai_enable_web_search').prop('checked', oai_settings.enable_web_search);
 
     $('#openai_reasoning_effort').val(oai_settings.reasoning_effort);
     $(`#openai_reasoning_effort option[value="${oai_settings.reasoning_effort}"]`).prop('selected', true);
@@ -3613,6 +3619,7 @@ async function saveOpenAIPreset(name, settings, triggerUi = true) {
         function_calling: settings.function_calling,
         show_thoughts: settings.show_thoughts,
         reasoning_effort: settings.reasoning_effort,
+        enable_web_search: settings.enable_web_search,
         seed: settings.seed,
         n: settings.n,
     };
@@ -5572,6 +5579,11 @@ export function initOpenAI() {
         saveSettingsDebounced();
     });
 
+    $('#openai_enable_web_search').on('input', function () {
+        oai_settings.enable_web_search = !!$(this).prop('checked');
+        saveSettingsDebounced();
+    });
+
     if (!CSS.supports('field-sizing', 'content')) {
         $(document).on('input', '#openai_settings .autoSetHeight', function () {
             resetScrollHeight($(this));
diff --git a/src/endpoints/backends/chat-completions.js b/src/endpoints/backends/chat-completions.js
index 857ecf9be..4899383dd 100644
--- a/src/endpoints/backends/chat-completions.js
+++ b/src/endpoints/backends/chat-completions.js
@@ -98,6 +98,21 @@ function getOpenRouterTransforms(request) {
     }
 }
 
+/**
+ * Gets OpenRouter plugins based on the request.
+ * @param {import('express').Request} request
+ * @returns {any[]} OpenRouter plugins
+ */
+function getOpenRouterPlugins(request) {
+    const plugins = [];
+
+    if (request.body.enable_web_search) {
+        plugins.push({ 'id': 'web' });
+    }
+
+    return plugins;
+}
+
 /**
  * Sends a request to Claude API.
  * @param {express.Request} request Express request
@@ -323,6 +338,7 @@ async function sendMakerSuiteRequest(request, response) {
 
     const model = String(request.body.model);
     const stream = Boolean(request.body.stream);
+    const enableWebSearch = Boolean(request.body.enable_web_search);
     const isThinking = model.includes('thinking');
 
     const generationConfig = {
@@ -348,6 +364,7 @@ async function sendMakerSuiteRequest(request, response) {
             model.startsWith('gemini-exp')
         ) && request.body.use_makersuite_sysprompt;
 
+        const tools = [];
         const prompt = convertGooglePrompt(request.body.messages, model, should_use_system_prompt, getPromptNames(request));
         let safetySettings = GEMINI_SAFETY;
 
@@ -361,6 +378,13 @@ async function sendMakerSuiteRequest(request, response) {
         }
         // Most of the other models allow for setting the threshold of filters, except for HARM_CATEGORY_CIVIC_INTEGRITY, to OFF.
 
+        if (enableWebSearch) {
+            const searchTool = model.includes('1.5') || model.includes('1.0')
+                ? ({ google_search_retrieval: {} })
+                : ({ google_search: {} });
+            tools.push(searchTool);
+        }
+
         let body = {
             contents: prompt.contents,
             safetySettings: safetySettings,
@@ -371,6 +395,10 @@ async function sendMakerSuiteRequest(request, response) {
             body.systemInstruction = prompt.system_instruction;
         }
 
+        if (tools.length) {
+            body.tools = tools;
+        }
+
         return body;
     }
 
@@ -1014,6 +1042,7 @@ router.post('/generate', jsonParser, function (request, response) {
         headers = { ...OPENROUTER_HEADERS };
         bodyParams = {
             'transforms': getOpenRouterTransforms(request),
+            'plugins': getOpenRouterPlugins(request),
             'include_reasoning': Boolean(request.body.include_reasoning),
         };