Add LLM token
This commit is contained in:
@@ -71,22 +71,48 @@ class OpenAICompatibleReranker(Reranker):
|
||||
return result
|
||||
|
||||
def _call_reranker(self, query: str, texts: list[str]) -> list[float]:
|
||||
"""Call the reranker API and return a score per text."""
|
||||
"""Call the reranker API and return a score per text.
|
||||
|
||||
Tries TEI format first (POST /rerank with model+texts), then falls back
|
||||
to Cohere/OpenAI format (POST /v1/rerank with model+documents).
|
||||
Both formats now include the model name, which most gateways require.
|
||||
"""
|
||||
headers = {"Content-Type": "application/json"}
|
||||
if self._api_key:
|
||||
headers["Authorization"] = f"Bearer {self._api_key}"
|
||||
|
||||
# Try TEI format first: POST /rerank
|
||||
payload = {"query": query, "texts": texts, "raw_scores": False, "return_text": False}
|
||||
# TEI format: POST /rerank — include model name (required by gateway proxies)
|
||||
payload = {
|
||||
"model": self._model,
|
||||
"query": query,
|
||||
"texts": texts,
|
||||
"raw_scores": False,
|
||||
"return_text": False,
|
||||
}
|
||||
url = f"{self._base_url}/rerank"
|
||||
resp = requests.post(url, json=payload, headers=headers, timeout=self._timeout)
|
||||
|
||||
if resp.status_code == 404:
|
||||
# Fall back to Cohere / OpenAI-style: POST /v1/rerank
|
||||
if resp.status_code in (404, 400):
|
||||
# Gateway returned an error — try Cohere/OpenAI-style format as fallback.
|
||||
logger.debug(
|
||||
"TEI rerank returned {} — trying Cohere format. Body: {}",
|
||||
resp.status_code,
|
||||
resp.text[:200],
|
||||
)
|
||||
payload_v1 = {"model": self._model, "query": query, "documents": texts}
|
||||
url = f"{self._base_url}/v1/rerank"
|
||||
resp = requests.post(url, json=payload_v1, headers=headers, timeout=self._timeout)
|
||||
|
||||
if not resp.ok:
|
||||
# Surface a clear error message so callers can log it meaningfully.
|
||||
try:
|
||||
err_body = resp.json()
|
||||
err_msg = err_body.get("error", {}).get("message", resp.text[:200])
|
||||
except Exception:
|
||||
err_msg = resp.text[:200]
|
||||
resp.raise_for_status() # raises HTTPError with status code
|
||||
raise ValueError(err_msg) # unreachable but satisfies type checker
|
||||
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
|
||||
|
||||
Reference in New Issue
Block a user