hunjing/backend/app/ai_provider/reranker/cohere_format.py

"""Cohere 兼容格式的 Reranker 适配器 — ``POST /v1/rerank``。

覆盖 Cohere、mixedbread、Jina AI 等采用 Cohere 兼容重排序端点的提供商。

请求::

    POST {base_url}/v1/rerank
    Authorization: Bearer {api_key}
    Content-Type: application/json

    {
        "model": "rerank-model",
        "query": "查询文本",
        "documents": ["文档1", "文档2"],
        "top_n": 3,
        "return_documents": false
    }

响应::

    {
        "results": [
            {"index": 0, "relevance_score": 0.99, "document": {"text": "文档..."}},
            {"index": 1, "relevance_score": 0.85, "document": {"text": "文档..."}}
        ],
        "meta": {"billed_units": {"search_units": 2}}
    }

另见:

    - Cohere Rerank API: https://docs.cohere.com/reference/rerank
    - Jina Reranker: https://api.jina.ai/redoc#tag/rerank
"""

from __future__ import annotations

from typing import Any

import httpx

from app.ai_provider.reranker.base import (
    RerankerAdapterBase,
    RerankerProtocol,
    RerankerServiceError,
    RerankRequest,
    RerankResponse,
    RerankResult,
)
from app.core.logging import logger


class CohereFormatRerankerAdapter(RerankerAdapterBase, RerankerProtocol):
    """Cohere 兼容重排序适配器，覆盖多数重排序服务提供商。"""

    _adapter_type = "reranker"
    _default_timeout = 60

    def __init__(
        self,
        *,
        base_url: str,
        api_key: str,
        model: str = "rerank-model",
        request_timeout_seconds: int = 60,
    ) -> None:
        self._base_url = base_url.rstrip("/")
        self._api_key = api_key
        self._model = model
        self._timeout = request_timeout_seconds
        self._client: httpx.AsyncClient | None = None

    # ── 客户端 ───────────────────────────────────────────────────────────

    @property
    def _http(self) -> httpx.AsyncClient:
        if self._client is None or self._client.is_closed:
            self._client = httpx.AsyncClient(timeout=httpx.Timeout(self._timeout))
        return self._client

    async def close(self) -> None:
        if self._client is not None and not self._client.is_closed:
            await self._client.aclose()

    # ── RerankerProtocol ────────────────────────────────────────────────

    async def rerank(self, request: RerankRequest) -> RerankResponse:
        """通过 ``POST /v1/rerank`` 对文档进行重排序。"""
        url = f"{self._base_url}/v1/rerank"

        payload: dict[str, Any] = {
            "model": request.model or self._model,
            "query": request.query,
            "documents": request.documents,
            "return_documents": request.return_documents,
        }
        if request.top_n is not None:
            payload["top_n"] = request.top_n
        if request.max_chunks_per_doc is not None:
            payload["max_chunks_per_doc"] = request.max_chunks_per_doc

        headers = {
            "Authorization": f"Bearer {self._api_key}",
            "Content-Type": "application/json",
        }

        logger.info(
            "reranker.request",
            model=payload["model"],
            query_length=len(request.query),
            document_count=len(request.documents),
            top_n=request.top_n,
            url=url,
        )

        try:
            response = await self._http.post(url, json=payload, headers=headers)
        except httpx.TimeoutException as exc:
            raise RerankerServiceError(
                status_code=504,
                code="reranker_timeout",
                message="Reranker request timed out.",
            ) from exc
        except httpx.RequestError as exc:
            raise RerankerServiceError(
                status_code=503,
                code="reranker_connection_error",
                message=f"Reranker connection error: {exc}",
            ) from exc

        if response.status_code != 200:
            body = response.text[:500]
            raise RerankerServiceError(
                status_code=response.status_code,
                code="reranker_api_error",
                message=f"Reranker API returned {response.status_code}: {body}",
            )

        result = self._parse_response(response.json())

        logger.info(
            "reranker.success",
            model=payload["model"],
            total_results=len(result.results),
        )
        return result

    # ── 辅助方法 ──────────────────────────────────────────────────────────

    @staticmethod
    def _parse_response(body: dict) -> RerankResponse:
        """从 Cohere 格式的重排序响应中提取结果。"""
        try:
            raw_results = body.get("results", [])
            results = []
            for item in raw_results:
                results.append(
                    RerankResult(
                        index=item.get("index", 0),
                        relevance_score=float(item.get("relevance_score", 0.0)),
                        text=item.get("document", {}).get("text")
                        if isinstance(item.get("document"), dict)
                        else None,
                    )
                )
            # 按相关性分数降序排列
            results.sort(key=lambda r: r.relevance_score, reverse=True)

            # 提取用量信息（不同提供商格式各异）
            usage: dict[str, int] | None = None
            meta = body.get("meta")
            if isinstance(meta, dict):
                billed = meta.get("billed_units")
                if isinstance(billed, dict):
                    usage = {k: int(v) for k, v in billed.items() if isinstance(v, (int, float))}

            return RerankResponse(
                results=results,
                model=body.get("model") or body.get("id", ""),
                usage=usage,
            )
        except Exception as exc:
            raise RerankerServiceError(
                status_code=500,
                code="reranker_response_parse_error",
                message=f"Failed to parse reranker response: {exc}",
            ) from exc