Prechádzať zdrojové kódy

enh: optionally add user headers external websearch

Co-Authored-By: Classic298 <27028174+Classic298@users.noreply.github.com>
Timothy Jaeryang Baek 3 mesiacov pred
rodič
commit
25c7f101f2

+ 20 - 5
backend/open_webui/retrieval/web/external.py

@@ -2,27 +2,42 @@ import logging
 from typing import Optional, List
 
 import requests
-from open_webui.retrieval.web.main import SearchResult, get_filtered_results
+
+from fastapi import Request
+
 from open_webui.env import SRC_LOG_LEVELS
 
+from open_webui.retrieval.web.main import SearchResult, get_filtered_results
+from open_webui.utils.headers import include_user_info_headers
+
+
 log = logging.getLogger(__name__)
 log.setLevel(SRC_LOG_LEVELS["RAG"])
 
 
 def search_external(
+    request: Request,
     external_url: str,
     external_api_key: str,
     query: str,
     count: int,
     filter_list: Optional[List[str]] = None,
+    user=None,
 ) -> List[SearchResult]:
     try:
+        headers = {
+            "User-Agent": "Open WebUI (https://github.com/open-webui/open-webui) RAG Bot",
+            "Authorization": f"Bearer {external_api_key}",
+        }
+        headers = include_user_info_headers(headers, user)
+
+        chat_id = getattr(request.state, "chat_id", None)
+        if chat_id:
+            headers["X-OpenWebUI-Chat-Id"] = str(chat_id)
+
         response = requests.post(
             external_url,
-            headers={
-                "User-Agent": "Open WebUI (https://github.com/open-webui/open-webui) RAG Bot",
-                "Authorization": f"Bearer {external_api_key}",
-            },
+            headers=headers,
             json={
                 "query": query,
                 "count": count,

+ 5 - 1
backend/open_webui/routers/retrieval.py

@@ -1811,7 +1811,9 @@ def process_web(
         )
 
 
-def search_web(request: Request, engine: str, query: str) -> list[SearchResult]:
+def search_web(
+    request: Request, engine: str, query: str, user=None
+) -> list[SearchResult]:
     """Search the web using a search engine and return the results as a list of SearchResult objects.
     Will look for a search engine API key in environment variables in the following order:
     - SEARXNG_QUERY_URL
@@ -2069,11 +2071,13 @@ def search_web(request: Request, engine: str, query: str) -> list[SearchResult]:
         )
     elif engine == "external":
         return search_external(
+            request,
             request.app.state.config.EXTERNAL_WEB_SEARCH_URL,
             request.app.state.config.EXTERNAL_WEB_SEARCH_API_KEY,
             query,
             request.app.state.config.WEB_SEARCH_RESULT_COUNT,
             request.app.state.config.WEB_SEARCH_DOMAIN_FILTER_LIST,
+            user=user,
         )
     else:
         raise Exception("No search engine API key found in environment variables")