Rerank 文档重排

import json
import math
import os
from openai import OpenAI, APIConnectionError, APIStatusError

query = "How can I reset my password?"
documents = [
    "Invoices are available in the billing dashboard.",
    "Use the password reset link on the sign-in page.",
    "You can change the interface language in settings.",
]
top_n = 2

with OpenAI(base_url=os.environ["API_BASE_URL"].rstrip("/"),
            api_key=os.environ["API_KEY"], timeout=60.0, max_retries=0) as client:
    try:
        # Generic HTTP helper: /rerank is a gateway extension, not an OpenAI API method.
        result = client.post("/rerank", cast_to=dict, body={
            "model": os.environ["API_MODEL"], "query": query, "documents": documents,
            "top_n": top_n, "return_documents": False,
        })
        if not isinstance(result, dict) or not isinstance(result.get("results"), list):
            raise ValueError("Expected a results array.")
        seen = set()
        matches = []
        for item in result["results"]:
            if not isinstance(item, dict):
                raise ValueError("Expected a result object.")
            index = item.get("index")
            score = item.get("relevance_score")
            if type(index) is not int or not 0 <= index < len(documents) or index in seen:
                raise ValueError("Invalid or duplicate document index.")
            if type(score) not in (int, float) or not math.isfinite(score):
                raise ValueError("Invalid relevance score.")
            seen.add(index)
            # Preserve application-owned text and IDs using the original input index.
            matches.append({"index": index, "score": score, "text": documents[index]})
        matches.sort(key=lambda item: item["score"], reverse=True)
        print(json.dumps(matches[:top_n], ensure_ascii=False, allow_nan=False))
    except APIStatusError as error:
        raise SystemExit(f"Request failed (HTTP {error.status_code}). Check rerank model access and parameters.")
    except APIConnectionError:
        raise SystemExit("Connection failed. Check the API address and network.")
    except (ValueError, TypeError, OverflowError):
        raise SystemExit("Invalid rerank response; check result indices and numeric scores.")