ailine/app/rag/reranker.py

"""
重排序器模块 (适配版)
使用远程 llama.cpp 服务 (兼容 OpenAI Rerank API) 替代本地 Cross-Encoder
"""
import requests
from typing import List
from langchain_core.documents import Document

class LLaMaCPPReranker:
    """使用远程 llama.cpp 服务对检索结果重排序。"""

    def __init__(self, 
                 base_url: str, 
                 api_key: str,
                 top_n: int = 5,
                 timeout: int = 60):
        """
        初始化远程重排序器
        
        Args:
            base_url: llama.cpp 服务的地址和端口，默认为环境变量 LLAMACPP_RERANKER_URL 或 "http://127.0.0.1:8083"。
            top_n: 返回前 N 个结果。
            api_key: API 密钥，默认为环境变量 LLAMACPP_API_KEY 或 "huang1998"。
            timeout: 请求超时时间（秒）。
        """
        self.base_url = base_url 
        self.api_key = api_key 
        self.top_n = top_n
        self.timeout = timeout
        self.endpoint = f"{self.base_url}/rerank"
        
    def compress_documents(
        self, documents: List[Document], query: str
    ) -> List[Document]:
        """
        对文档进行重排序
        
        Args:
            documents: 待排序的文档列表
            query: 查询字符串
        
        Returns:
            排序后的文档列表
        """
        if not documents:
            return []

        # 准备请求体
        # 根据 llama.cpp 的 OpenAI 兼容性，文档是一个字符串列表
        payload = {
            "model": "bge-reranker-v2-m3",
            "query": query,
            "documents": [doc.page_content for doc in documents],
            "top_n": self.top_n
        }
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }

        try:
            response = requests.post(self.endpoint, json=payload, headers=headers, timeout=self.timeout)
            response.raise_for_status()  # 检查请求是否成功
            results = response.json()
            
            # 解析返回结果
            # 返回格式: {"results": [{"index": 0, "document": "...", "relevance_score": 0.8}, ...]}
            # 按相关性得分降序排列
            sorted_indices = [item["index"] for item in results["results"]]
            sorted_docs = [documents[idx] for idx in sorted_indices]
            return sorted_docs
            
        except Exception as e:
            print(f"警告: 远程重排序过程出错，将使用原始排序。错误: {e}")
            return documents[:self.top_n]