refac/enh: forward user info header to reranker

backend/open_webui/retrieval/models/external.py

@@ -1,8 +1,10 @@
 import logging
 import requests
 from typing import Optional, List, Tuple
+from urllib.parse import quote
 
-from open_webui.env import SRC_LOG_LEVELS
+
+from open_webui.env import ENABLE_FORWARD_USER_INFO_HEADERS, SRC_LOG_LEVELS
 from open_webui.retrieval.models.base_reranker import BaseReranker
 
 
@@ -21,7 +23,9 @@ class ExternalReranker(BaseReranker):
         self.url = url
         self.model = model
 
-    def predict(self, sentences: List[Tuple[str, str]]) -> Optional[List[float]]:
+    def predict(
+        self, sentences: List[Tuple[str, str]], user=None
+    ) -> Optional[List[float]]:
         query = sentences[0][0]
         docs = [i[1] for i in sentences]
 
@@ -41,6 +45,16 @@ class ExternalReranker(BaseReranker):
                 headers={
                     "Content-Type": "application/json",
                     "Authorization": f"Bearer {self.api_key}",
+                    **(
+                        {
+                            "X-OpenWebUI-User-Name": quote(user.name, safe=" "),
+                            "X-OpenWebUI-User-Id": user.id,
+                            "X-OpenWebUI-User-Email": user.email,
+                            "X-OpenWebUI-User-Role": user.role,
+                        }
+                        if ENABLE_FORWARD_USER_INFO_HEADERS and user
+                        else {}
+                    ),
                 },
                 json=payload,
             )

backend/open_webui/retrieval/utils.py

@@ -445,6 +445,15 @@ def get_embedding_function(
         raise ValueError(f"Unknown embedding engine: {embedding_engine}")
 
 
+def get_reranking_function(reranking_engine, reranking_model, reranking_function):
+    if reranking_engine == "external":
+        return lambda sentences, user=None: reranking_function.predict(
+            sentences, user=user
+        )
+    else:
+        return lambda sentences, user=None: reranking_function.predict(sentences)
+
+
 def get_sources_from_items(
     request,
     items,
@@ -925,7 +934,7 @@ class RerankCompressor(BaseDocumentCompressor):
         reranking = self.reranking_function is not None
 
         if reranking:
-            scores = self.reranking_function.predict(
+            scores = self.reranking_function(
                 [(query, doc.page_content) for doc in documents]
             )
         else: