Browse Source

enh: query caching

Co-Authored-By: Jacob Leksan <63938553+jmleksan@users.noreply.github.com>
Timothy Jaeryang Baek 1 month ago
parent
commit
31485835a7

+ 2 - 0
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
 ####################################

+ 4 - 0
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,

+ 4 - 0
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]