mirror of
https://github.com/open-webui/open-webui.git
synced 2026-11-15 18:55:07 +03:00
Classic298
8d7b71f6cc
The external reranker, Exa and Ollama web search wrote the full user query to the server log at info level, and reusing generated search queries logged them as well. Those lines no longer include the query. A failing external reranker also logged the provider URL in its error message along with a stack trace; it now logs only the error type. Co-authored-by: manooll <16244686+manooll@users.noreply.github.com>
69 lines
2.0 KiB
Python
69 lines
2.0 KiB
Python
import logging
|
|
from typing import List, Optional, Tuple
|
|
from urllib.parse import quote
|
|
|
|
import requests
|
|
from open_webui.env import ENABLE_FORWARD_USER_INFO_HEADERS, REQUESTS_VERIFY
|
|
from open_webui.retrieval.models.base_reranker import BaseReranker
|
|
from open_webui.utils.headers import include_user_info_headers
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
class ExternalReranker(BaseReranker):
|
|
def __init__(
|
|
self,
|
|
api_key: str,
|
|
url: str = 'http://localhost:8080/v1/rerank',
|
|
model: str = 'reranker',
|
|
timeout: Optional[int] = None,
|
|
):
|
|
self.api_key = api_key
|
|
self.url = url
|
|
self.model = model
|
|
self.timeout = timeout
|
|
|
|
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),
|
|
}
|
|
|
|
try:
|
|
log.info('ExternalReranker:predict:model %s', self.model)
|
|
|
|
headers = {
|
|
'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}',
|
|
headers=headers,
|
|
json=payload,
|
|
timeout=self.timeout,
|
|
verify=REQUESTS_VERIFY,
|
|
)
|
|
|
|
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]
|
|
else:
|
|
log.error('No results found in external reranking response')
|
|
return None
|
|
|
|
except Exception as e:
|
|
log.error('Error in external reranking: %s', type(e).__name__)
|
|
return None
|