From 31485835a7f3b510a2ff930ce9afcafdbc770b0a Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Wed, 27 Aug 2025 03:07:21 +0400
Subject: [PATCH] enh: query caching

Co-Authored-By: Jacob Leksan <63938553+jmleksan@users.noreply.github.com>
---
 backend/open_webui/env.py              | 2 ++
 backend/open_webui/routers/tasks.py    | 4 ++++
 backend/open_webui/utils/middleware.py | 4 ++++
 3 files changed, 10 insertions(+)

diff --git a/backend/open_webui/env.py b/backend/open_webui/env.py
index 937b06794f..4762a66633 100644
--- a/backend/open_webui/env.py
+++ b/backend/open_webui/env.py
@@ -362,6 +362,8 @@ ENABLE_REALTIME_CHAT_SAVE = (
     os.environ.get("ENABLE_REALTIME_CHAT_SAVE", "False").lower() == "true"
 )
 
+ENABLE_QUERIES_CACHE = os.environ.get("ENABLE_QUERIES_CACHE", "False").lower() == "true"
+
 ####################################
 # REDIS
 ####################################
diff --git a/backend/open_webui/routers/tasks.py b/backend/open_webui/routers/tasks.py
index e49602094f..7585466f69 100644
--- a/backend/open_webui/routers/tasks.py
+++ b/backend/open_webui/routers/tasks.py
@@ -470,6 +470,10 @@ async def generate_queries(
                 detail=f"Query generation is disabled",
             )
 
+    if getattr(request.state, "cached_queries", None):
+        log.info(f"Reusing cached queries: {request.state.cached_queries}")
+        return request.state.cached_queries
+
     if getattr(request.state, "direct", False) and hasattr(request.state, "model"):
         models = {
             request.state.model["id"]: request.state.model,
diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py
index 314885f0b8..66d4ad6286 100644
--- a/backend/open_webui/utils/middleware.py
+++ b/backend/open_webui/utils/middleware.py
@@ -101,6 +101,7 @@ from open_webui.env import (
     CHAT_RESPONSE_MAX_TOOL_CALL_RETRIES,
     BYPASS_MODEL_ACCESS_CONTROL,
     ENABLE_REALTIME_CHAT_SAVE,
+    ENABLE_QUERIES_CACHE,
 )
 from open_webui.constants import TASKS
 
@@ -391,6 +392,9 @@ async def chat_web_search_handler(
         except Exception as e:
             queries = [response]
 
+        if ENABLE_QUERIES_CACHE:
+            request.state.cached_queries = queries
+
     except Exception as e:
         log.exception(e)
         queries = [user_message]