Spaces:
Build error
Build error
File size: 2,151 Bytes
87a665c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | import logging
import requests
from typing import Optional, List, Tuple
from urllib.parse import quote
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(f'ExternalReranker:predict:model {self.model}')
log.info(f'ExternalReranker:predict:query {query}')
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.exception(f'Error in external reranking: {e}')
return None
|