From 9ca47275730487351deca40fd7d31dc8600d92ca Mon Sep 17 00:00:00 2001
From: Timothy Jaeryang Baek <tim@openwebui.com>
Date: Tue, 18 Feb 2025 21:29:27 -0800
Subject: [PATCH] enh: web search behaviour

---
 backend/open_webui/utils/middleware.py | 123 +++++++++++++------------
 1 file changed, 64 insertions(+), 59 deletions(-)

diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py
index b624f2a34..baa4d49a4 100644
--- a/backend/open_webui/utils/middleware.py
+++ b/backend/open_webui/utils/middleware.py
@@ -321,89 +321,94 @@ async def chat_web_search_handler(
         )
         return form_data
 
-    searchQuery = queries[0]
+    all_results = []
 
-    await event_emitter(
-        {
-            "type": "status",
-            "data": {
-                "action": "web_search",
-                "description": 'Searching "{{searchQuery}}"',
-                "query": searchQuery,
-                "done": False,
-            },
-        }
-    )
-
-    try:
-
-        results = await process_web_search(
-            request,
-            SearchForm(
-                **{
+    for searchQuery in queries:
+        await event_emitter(
+            {
+                "type": "status",
+                "data": {
+                    "action": "web_search",
+                    "description": 'Searching "{{searchQuery}}"',
                     "query": searchQuery,
-                }
-            ),
-            user=user,
+                    "done": False,
+                },
+            }
         )
 
-        if results:
-            await event_emitter(
-                {
-                    "type": "status",
-                    "data": {
-                        "action": "web_search",
-                        "description": "Searched {{count}} sites",
+        try:
+            results = await process_web_search(
+                request,
+                SearchForm(
+                    **{
                         "query": searchQuery,
-                        "urls": results["filenames"],
-                        "done": True,
-                    },
-                }
+                    }
+                ),
+                user=user,
             )
 
-            files = form_data.get("files", [])
+            if results:
+                all_results.append(results)
+                files = form_data.get("files", [])
 
-            if request.app.state.config.RAG_WEB_SEARCH_FULL_CONTEXT:
-                files.append(
-                    {
-                        "docs": results.get("docs", []),
-                        "name": searchQuery,
-                        "type": "web_search_docs",
-                        "urls": results["filenames"],
-                    }
-                )
-            else:
-                files.append(
-                    {
-                        "collection_name": results["collection_name"],
-                        "name": searchQuery,
-                        "type": "web_search_results",
-                        "urls": results["filenames"],
-                    }
-                )
-            form_data["files"] = files
-        else:
+                if request.app.state.config.RAG_WEB_SEARCH_FULL_CONTEXT:
+                    files.append(
+                        {
+                            "docs": results.get("docs", []),
+                            "name": searchQuery,
+                            "type": "web_search_docs",
+                            "urls": results["filenames"],
+                        }
+                    )
+                else:
+                    files.append(
+                        {
+                            "collection_name": results["collection_name"],
+                            "name": searchQuery,
+                            "type": "web_search_results",
+                            "urls": results["filenames"],
+                        }
+                    )
+                form_data["files"] = files
+        except Exception as e:
+            log.exception(e)
             await event_emitter(
                 {
                     "type": "status",
                     "data": {
                         "action": "web_search",
-                        "description": "No search results found",
+                        "description": 'Error searching "{{searchQuery}}"',
                         "query": searchQuery,
                         "done": True,
                         "error": True,
                     },
                 }
             )
-    except Exception as e:
-        log.exception(e)
+
+    if all_results:
+        urls = []
+        for results in all_results:
+            if "filenames" in results:
+                urls.extend(results["filenames"])
+
         await event_emitter(
             {
                 "type": "status",
                 "data": {
                     "action": "web_search",
-                    "description": 'Error searching "{{searchQuery}}"',
-                    "query": searchQuery,
+                    "description": "Searched {{count}} sites",
+                    "urls": urls,
+                    "done": True,
+                },
+            }
+        )
+    else:
+        await event_emitter(
+            {
+                "type": "status",
+                "data": {
+                    "action": "web_search",
+                    "description": "No search results found",
                     "done": True,
                     "error": True,
                 },