refac

backend/open_webui/retrieval/models/external.py

@@ -15,8 +15,8 @@ class ExternalReranker(BaseReranker):
     def __init__(
         self,
         api_key: str,
-        url: str = "http://localhost:8080/v1/rerank",
-        model: str = "reranker",
+        url: str = 'http://localhost:8080/v1/rerank',
+        model: str = 'reranker',
         timeout: Optional[int] = None,
     ):
         self.api_key = api_key
@@ -24,33 +24,31 @@ class ExternalReranker(BaseReranker):
         self.model = model
         self.timeout = timeout
 
-    def predict(
-        self, sentences: List[Tuple[str, str]], user=None
-    ) -> 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]
 
         payload = {
-            "model": self.model,
-            "query": query,
-            "documents": docs,
-            "top_n": len(docs),
+            'model': self.model,
+            'query': query,
+            'documents': docs,
+            'top_n': len(docs),
         }
 
         try:
-            log.info(f"ExternalReranker:predict:model {self.model}")
-            log.info(f"ExternalReranker:predict:query {query}")
+            log.info(f'ExternalReranker:predict:model {self.model}')
+            log.info(f'ExternalReranker:predict:query {query}')
 
             headers = {
-                "Content-Type": "application/json",
-                "Authorization": f"Bearer {self.api_key}",
+                'Content-Type': 'application/json',
+                'Authorization': f'Bearer {self.api_key}',
             }
 
             if ENABLE_FORWARD_USER_INFO_HEADERS and user:
                 headers = include_user_info_headers(headers, user)
 
             r = requests.post(
-                f"{self.url}",
+                f'{self.url}',
                 headers=headers,
                 json=payload,
                 timeout=self.timeout,
@@ -60,13 +58,13 @@ class ExternalReranker(BaseReranker):
             r.raise_for_status()
             data = r.json()
 
-            if "results" in data:
-                sorted_results = sorted(data["results"], key=lambda x: x["index"])
-                return [result["relevance_score"] for result in sorted_results]
+            if 'results' in data:
+                sorted_results = sorted(data['results'], key=lambda x: x['index'])
+                return [result['relevance_score'] for result in sorted_results]
             else:
-                log.error("No results found in external reranking response")
+                log.error('No results found in external reranking response')
                 return None
 
         except Exception as e:
-            log.exception(f"Error in external reranking: {e}")
+            log.exception(f'Error in external reranking: {e}')
             return None