LQAgentPlatform/foundation/ai/models/rerank_model.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
重排序执行模块
用于调用BGE重排序模型进行文档重排序
"""
import json
import requests
from typing import List, Dict, Any
from foundation.infrastructure.config.config import config_handler
from foundation.observability.logger.loggering import server_logger


class LqReranker:
    """
    重排序执行器
    """

    def __init__(self):
        self.api_url = config_handler.get('rerank_model', 'BGE_RERANKER_SERVER_RUL')
        self.model = config_handler.get('rerank_model', 'BGE_RERANKER_MODEL_ID')
        # 确保top_k是整数类型，避免切片错误
        self.top_k = int(config_handler.get('rerank_model', 'BGE_RERANKER_TOP_N', 5))
        
    def bge_rerank(self,query: str, candidates: List[str],top_k :int = None) -> List[Dict[str, Any]]:
        """
        执行重排序的全局函数

        Args:
            query: 查询文本
            candidates: 候选文档列表
            top_k: 调用时chaurnum参数，默认为None


        Returns:
            List[Dict]: 重排序后的结果列表
        """
        try:
            # self.top_k 是config.ini生产环境中实际使用的重排序数量，bge_rerank中的top_k，用于开发环境中快速效果调试
            if not top_k:# 如果开发top_k未指定，则使用配置文件中的top_k
                top_k = self.top_k
            

            server_logger.info(f"开始执行重排序，查询:, 候选文档数量: {len(candidates)}")

            # 构建重排序请求
            rerank_request = {
                "model": "bge-reranker-v2-m3",
                "query": query,
                "candidates": candidates
            }

            # 直接调用重排序API
            url = self.api_url
            headers = {
                "Content-Type": "application/json"
            }

            server_logger.debug(f"调用重排序API: {url}")
            server_logger.debug(f"请求数据: {json.dumps(rerank_request, ensure_ascii=False)}")

            response = requests.post(url, headers=headers, json=rerank_request, timeout=30)

            if response.status_code == 200:
                result = response.json()
                server_logger.debug(f"API响应: {json.dumps(result, ensure_ascii=False)}")

                if "results" in result:
                    return result["results"][:top_k]
                else:
                    server_logger.warning(f"API响应格式异常: {result}")
                    return []
            else:
                server_logger.error(f"API调用失败，状态码: {response.status_code}, 响应: {response.text}")
                return []

        except Exception as e:
            server_logger.error(f"执行重排序失败: {str(e)}")
            # 返回原始顺序作为fallback
            return [{"text": doc, "score": "0.0"} for doc in candidates[:top_k]]

rerank_model = LqReranker()