mildoc/mildoc_wxkf/rerank_service.py

"""
重排序服务模块

支持多个重排序服务提供商：
- 阿里百炼平台 (dashscope)
- 硅基流动平台 (siliconflow)

作者：开发工程师
日期：2025年01月
"""

import logging
import requests
import json
from enum import Enum
from typing import List, Optional, Dict, Any
from pydantic import BaseModel
from config import Config

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class RerankProvider(Enum):
    """重排序服务提供商"""
    DASHSCOPE = "dashscope"  # 阿里百炼
    SILICONFLOW = "siliconflow"  # 硅基流动


class RerankDocument(BaseModel):
    """重排序文档模型"""
    index: int  # 原始文档索引
    content: str  # 文档内容
    relevance_score: float  # 相关性分数
    metadata: Optional[Dict[str, Any]] = None  # 元数据


class RerankResponse(BaseModel):
    """重排序响应模型"""
    documents: List[RerankDocument]  # 重排序后的文档
    success: bool = True
    error_message: Optional[str] = None


class RerankService:
    """重排序服务类
    
    支持多个重排序服务提供商，通过统一接口提供文档重排序功能
    """
    
    def __init__(self, provider: RerankProvider, api_key: str, model_name: str, endpoint: Optional[str] = None):
        """初始化重排序服务
        
        Args:
            provider: 服务提供商
            api_key: API密钥
            model_name: 模型名称
            endpoint: 自定义API端点（可选）
        """
        self.provider = provider
        self.api_key = api_key
        self.model_name = model_name
        self.endpoint = endpoint
        logger.info(f"重排序服务初始化完成: provider={provider.value}, model={model_name}, endpoint={self.endpoint}")
    
    def rerank_documents(self, query: str, documents: List[str], top_n: Optional[int] = None) -> RerankResponse:
        """重排序文档
        
        Args:
            query: 查询文本
            documents: 文档列表
            top_n: 返回前N个文档，默认返回全部
            
        Returns:
            RerankResponse: 重排序结果
        """
        try:
            if not query or not documents:
                logger.warning("查询或文档列表为空")
                return RerankResponse(
                    documents=[],
                    success=False,
                    error_message="查询或文档列表为空"
                )
            
            logger.info(f"🔄 开始重排序: 查询='{query[:50]}...', 文档数量={len(documents)}, top_n={top_n}")
            
            if self.provider == RerankProvider.DASHSCOPE:
                response = self._rerank_dashscope(query, documents, top_n)
            elif self.provider == RerankProvider.SILICONFLOW:
                response = self._rerank_siliconflow(query, documents, top_n)
            else:
                raise ValueError(f"不支持的重排序提供商: {self.provider}")
            
            if response.success:
                logger.info(f"✅ 重排序完成: 返回{len(response.documents)}个文档")
                for i, doc in enumerate(response.documents): 
                    logger.info(f"  #{i+1}: 相关性={doc.relevance_score:.4f}, 内容='{doc.content[:50]}...'")
            
            return response
                
        except Exception as e:
            logger.error(f"❌ 重排序失败: {e}")
            return RerankResponse(
                documents=[],
                success=False,
                error_message=str(e)
            )
    
    def _rerank_dashscope(self, query: str, documents: List[str], top_n: Optional[int] = None) -> RerankResponse:
        """阿里百炼平台重排序"""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        data = {
            "model": self.model_name,
            "input": {
                "query": query,
                "documents": documents
            },
            "parameters": {
                "return_documents": True,  # 显式设置返回文档内容
                "top_n": top_n or len(documents)
            }
        }
        
        logger.debug(f"🌐 发送请求到百炼平台: {self.endpoint}")
        response = requests.post(self.endpoint, headers=headers, json=data, timeout=30)
        response.raise_for_status()
        
        result = response.json()
        logger.debug(f"📨 百炼平台响应: {json.dumps(result, ensure_ascii=False)[:200]}...")
        
        # 解析响应
        rerank_docs = []
        if "output" in result and "results" in result["output"]:
            for item in result["output"]["results"]:
                # 百炼平台返回格式: document.text
                document = item.get("document", {})
                content = document.get("text", "") if isinstance(document, dict) else str(document)
                
                rerank_doc = RerankDocument(
                    index=item.get("index", 0),
                    content=content,
                    relevance_score=float(item.get("relevance_score", 0.0))
                )
                rerank_docs.append(rerank_doc)
        else:
            logger.warning(f"百炼平台响应格式异常: {result}")
            return RerankResponse(
                documents=[],
                success=False,
                error_message="响应格式异常"
            )
        
        return RerankResponse(documents=rerank_docs)
    
    def _rerank_siliconflow(self, query: str, documents: List[str], top_n: Optional[int] = None) -> RerankResponse:
        """硅基流动平台重排序"""
        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }
        
        data = {
            "model": self.model_name,
            "query": query,
            "documents": documents,
            "return_documents": True  # 硅基流动需要明确指定返回文档内容
        }
        
        # 硅基流动支持top_n参数
        if top_n is not None:
            data["top_n"] = top_n
        
        logger.debug(f"🌐 发送请求到硅基流动: {self.endpoint}")
        response = requests.post(self.endpoint, headers=headers, json=data, timeout=30)
        response.raise_for_status()
        
        result = response.json()
        logger.debug(f"📨 硅基流动响应: {json.dumps(result, ensure_ascii=False)[:200]}...")
        
        # 解析响应
        rerank_docs = []
        if "results" in result:
            for item in result["results"]:
                # 硅基流动平台返回格式: document.text
                document = item.get("document", {})
                content = document.get("text", "") if isinstance(document, dict) else str(document)
                
                rerank_doc = RerankDocument(
                    index=item.get("index", 0),
                    content=content,
                    relevance_score=float(item.get("relevance_score", 0.0))
                )
                rerank_docs.append(rerank_doc)
        else:
            logger.warning(f"硅基流动响应格式异常: {result}")
            return RerankResponse(
                documents=[],
                success=False,
                error_message="响应格式异常"
            )
        
        return RerankResponse(documents=rerank_docs)
    
    def health_check(self) -> Dict[str, Any]:
        """健康检查"""
        status = {
            "service": "RerankService",
            "provider": self.provider.value,
            "model": self.model_name,
            "endpoint": self.endpoint,
            "status": "unknown"
        }
        
        try:
            # 简单的健康检查：用最少的文档进行测试
            test_response = self.rerank_documents(
                query="测试",
                documents=["这是一个测试文档"],
                top_n=1
            )
            
            if test_response.success:
                status["status"] = "healthy"
                status["test_result"] = "ok"
            else:
                status["status"] = "error"
                status["error"] = test_response.error_message
                
        except Exception as e:
            status["status"] = "error"
            status["error"] = str(e)
        
        return status


def create_rerank_service() -> Optional[RerankService]:
    """创建重排序服务实例（工厂函数）
    
    Returns:
        RerankService: 重排序服务实例，失败时返回None
    """
    try:
        # 检查配置
        if not Config.RERANK_PROVIDER:
            logger.info("未配置RERANK_PROVIDER，跳过重排序服务初始化")
            return None
        
        if not Config.RERANK_API_KEY:
            logger.warning("缺少RERANK_API_KEY配置")
            return None
            
        if not Config.RERANK_MODEL_NAME:
            logger.warning("缺少RERANK_MODEL_NAME配置")
            return None
        
        # 创建提供商枚举
        try:
            provider = RerankProvider(Config.RERANK_PROVIDER.lower())
        except ValueError:
            logger.error(f"不支持的重排序提供商: {Config.RERANK_PROVIDER}")
            return None
        
        # 创建服务实例
        service = RerankService(
            provider=provider,
            api_key=Config.RERANK_API_KEY,
            model_name=Config.RERANK_MODEL_NAME,
            endpoint=Config.RERANK_ENDPOINT
        )
        
        logger.info(f"重排序服务创建成功: {provider.value}")
        return service
        
    except Exception as e:
        logger.error(f"创建重排序服务失败: {e}")
        return None


# 全局重排序服务实例
_rerank_service_instance = None

def get_rerank_service() -> Optional[RerankService]:
    """获取重排序服务实例（单例模式）"""
    global _rerank_service_instance
    
    if _rerank_service_instance is None:
        _rerank_service_instance = create_rerank_service()
    
    return _rerank_service_instance


if __name__ == "__main__":
    # 测试代码
    rerank_service = get_rerank_service()
    if rerank_service:
        # 健康检查
        health = rerank_service.health_check()
        print(f"健康检查: {health}")
        
        # 测试重排序
        test_docs = [
            "苹果是一种水果，含有丰富的维生素",
            "量子计算是计算科学的前沿领域，苹果也在做这方面的研究",
            "苹果是一家科技企业，主要生产手机、电脑等产品"
        ]
        
        result = rerank_service.rerank_documents(
            query="苹果公司",
            documents=test_docs,
            top_n=3
        )
        
        if result.success:
            print("\n重排序结果:")
            for i, doc in enumerate(result.documents):
                print(f"{i+1}. [分数: {doc.relevance_score:.4f}] {doc.content}")
        else:
            print(f"重排序失败: {result.error_message}")
    else:
        print("重排序服务未配置或初始化失败")