feat(cloudflare): add rerank support

Co-authored-by: Codex <codex@openai.com>

README.md

@@ -292,7 +292,7 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse
 | [Bytez (`bytez`)](https://docs.litellm.ai/docs/providers/bytez) | ✅ | ✅ | ✅ |  |  |  |  |  |  |  |
 | [Cerebras (`cerebras`)](https://docs.litellm.ai/docs/providers/cerebras) | ✅ | ✅ | ✅ |  |  |  |  |  |  |  |
 | [Clarifai (`clarifai`)](https://docs.litellm.ai/docs/providers/clarifai) | ✅ | ✅ | ✅ |  |  |  |  |  |  |  |
-| [Cloudflare AI Workers (`cloudflare`)](https://docs.litellm.ai/docs/providers/cloudflare_workers) | ✅ | ✅ | ✅ | ✅ |  |  |  |  |  |  |
+| [Cloudflare AI Workers (`cloudflare`)](https://docs.litellm.ai/docs/providers/cloudflare_workers) | ✅ | ✅ | ✅ | ✅ |  |  |  |  |  | ✅ |
 | [Codestral (`codestral`)](https://docs.litellm.ai/docs/providers/codestral) | ✅ | ✅ | ✅ |  |  |  |  |  |  |  |
 | [Cognition (`cognition`)](https://docs.litellm.ai/docs/providers/cognition) | ✅ | ✅ | ✅ |  |  |  |  |  |  |  |
 | [Cohere (`cohere`)](https://docs.litellm.ai/docs/providers/cohere) | ✅ | ✅ | ✅ | ✅ |  |  |  |  |  | ✅ |

litellm/__init__.py

@@ -1646,6 +1646,9 @@ if TYPE_CHECKING:
     from .llms.cloudflare.embedding.transformation import (
         CloudflareEmbeddingConfig as CloudflareEmbeddingConfig,
     )
+    from .llms.cloudflare.rerank.transformation import (
+        CloudflareRerankConfig as CloudflareRerankConfig,
+    )
     from .llms.novita.chat.transformation import NovitaConfig as NovitaConfig
     from .llms.petals.completion.transformation import PetalsConfig as PetalsConfig
     from .llms.ollama.chat.transformation import OllamaChatConfig as OllamaChatConfig

litellm/_lazy_imports_registry.py

@@ -165,6 +165,7 @@ LLM_CONFIG_NAMES: Final = (
     "TogetherAITextCompletionConfig",
     "CloudflareChatConfig",
     "CloudflareEmbeddingConfig",
+    "CloudflareRerankConfig",
     "NovitaConfig",
     "PetalsConfig",
     "OllamaChatConfig",
@@ -719,6 +720,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
         ".llms.cloudflare.embedding.transformation",
         "CloudflareEmbeddingConfig",
     ),
+    "CloudflareRerankConfig": (
+        ".llms.cloudflare.rerank.transformation",
+        "CloudflareRerankConfig",
+    ),
     "NovitaConfig": (".llms.novita.chat.transformation", "NovitaConfig"),
     "PetalsConfig": (".llms.petals.completion.transformation", "PetalsConfig"),
     "OllamaChatConfig": (".llms.ollama.chat.transformation", "OllamaChatConfig"),

litellm/llms/cloudflare/rerank/transformation.py

@@ -0,0 +1,284 @@
+import json
+from collections.abc import Mapping, Sequence
+from typing import Union
+
+import httpx
+from typing_extensions import NotRequired, ReadOnly, TypedDict
+
+from litellm._uuid import uuid
+from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
+from litellm.litellm_core_utils.url_utils import encode_url_path_segments
+from litellm.llms.base_llm.chat.transformation import BaseLLMException
+from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig
+from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str
+from litellm.types.rerank import RerankResponse, RerankResponseResult
+
+from ..chat.transformation import CloudflareError
+
+
+class CloudflareRerankContext(TypedDict):
+    text: ReadOnly[str]
+
+
+class CloudflareRerankRequest(TypedDict):
+    query: ReadOnly[str]
+    contexts: ReadOnly[Sequence[Mapping[str, str]]]
+    top_k: NotRequired[ReadOnly[object]]
+
+
+class CohereRerankParams(TypedDict):
+    query: ReadOnly[str]
+    documents: ReadOnly[Sequence[Union[str, Mapping[str, object]]]]
+    top_n: NotRequired[ReadOnly[int]]
+    return_documents: NotRequired[ReadOnly[bool]]
+
+
+class CloudflareHeaders(TypedDict):
+    Authorization: str
+    accept: str
+
+
+class LoggingAdditionalArgs(TypedDict):
+    complete_input_dict: ReadOnly[Mapping[str, object]]
+
+
+class EmptyRequestData(TypedDict, total=False):
+    pass
+
+
+EMPTY_REQUEST_DATA: Mapping[str, object] = EmptyRequestData()
+
+
+class CloudflareRerankConfig(BaseRerankConfig):
+    @staticmethod
+    def _default_api_base() -> str:
+        account_id = normalize_nonempty_secret_str(get_secret_str("CLOUDFLARE_ACCOUNT_ID"))
+        if account_id is None:
+            raise ValueError("Missing CLOUDFLARE_ACCOUNT_ID - set CLOUDFLARE_ACCOUNT_ID or pass api_base explicitly")
+        return f"https://api.cloudflare.com/client/v4/accounts/{account_id}/ai/run"
+
+    @staticmethod
+    def _document_to_context(
+        document: object,
+    ) -> Mapping[str, str]:
+        if isinstance(document, str):
+            return CloudflareRerankContext(text=document)
+        if not isinstance(document, Mapping):
+            raise TypeError("Cloudflare rerank documents must be strings or dictionaries")
+
+        text = document.get("text")
+        return CloudflareRerankContext(text=text if isinstance(text, str) else json.dumps(document))
+
+    @staticmethod
+    def _response_items(
+        response_json: Mapping[str, object],
+        status_code: int,
+    ) -> tuple[object, ...]:
+        if response_json.get("success") is False:
+            raise CloudflareError(
+                status_code=status_code,
+                message=str(response_json.get("errors") or response_json),
+            )
+
+        result = response_json.get("result", response_json)
+        response = result.get("response") if isinstance(result, Mapping) else None
+        if not isinstance(response, list):
+            raise CloudflareError(
+                status_code=status_code,
+                message=f"No response in Cloudflare rerank result: {response_json}",
+            )
+        return tuple(response)
+
+    @staticmethod
+    def _transform_response_item(
+        item: object,
+        documents: Sequence[object],
+        return_documents: bool,
+    ) -> RerankResponseResult:
+        if not isinstance(item, Mapping):
+            raise TypeError("Invalid item in Cloudflare rerank response")
+
+        index = item.get("id")
+        score = item.get("score")
+        if not isinstance(index, int) or not isinstance(score, (int, float)):
+            raise TypeError("Invalid item in Cloudflare rerank response")
+
+        if return_documents and index < len(documents):
+            document = CloudflareRerankConfig._document_to_context(documents[index])
+            return RerankResponseResult(
+                index=index,
+                relevance_score=float(score),
+                document=document,
+            )
+        return RerankResponseResult(index=index, relevance_score=float(score))
+
+    def validate_environment(
+        self,
+        headers: Mapping[str, object],
+        model: str,
+        api_key: str | None = None,
+        optional_params: Mapping[str, object] | None = None,
+    ) -> Mapping[str, object]:
+        api_key = api_key or get_secret_str("CLOUDFLARE_API_KEY")
+        if api_key is None:
+            raise ValueError("Missing Cloudflare API Key - set CLOUDFLARE_API_KEY or pass api_key explicitly")
+        cloudflare_headers = CloudflareHeaders(
+            Authorization=f"Bearer {api_key}",
+            accept="application/json",
+        )
+        cloudflare_headers["content-type"] = "application/json"
+        cloudflare_headers.update(headers)
+        return cloudflare_headers
+
+    def get_complete_url(
+        self,
+        api_base: str | None,
+        model: str,
+        optional_params: Mapping[str, object] | None = None,
+    ) -> str:
+        cleaned = (api_base or self._default_api_base()).rstrip("/")
+        encoded_model = encode_url_path_segments(model, field_name="model")
+        if cleaned.endswith(f"/{encoded_model}"):
+            return cleaned
+        if cleaned.endswith("/ai/v1"):
+            return f"{cleaned[: -len('/ai/v1')]}/ai/run/{encoded_model}"
+        if cleaned.endswith("/ai/run"):
+            return f"{cleaned}/{encoded_model}"
+        return f"{cleaned}/ai/run/{encoded_model}"
+
+    def get_supported_cohere_rerank_params(
+        self,
+        model: str,
+    ) -> Sequence[str]:
+        return (
+            "query",
+            "documents",
+            "top_n",
+            "return_documents",
+        )
+
+    def map_cohere_rerank_params(
+        self,
+        non_default_params: Mapping[str, object],
+        model: str,
+        drop_params: bool,
+        query: str,
+        documents: Sequence[
+            Union[
+                str,
+                Mapping[str, object],
+            ]
+        ],
+        custom_llm_provider: str | None = None,
+        top_n: int | None = None,
+        rank_fields: Sequence[str] | None = None,
+        return_documents: bool | None = True,
+        max_chunks_per_doc: int | None = None,
+        max_tokens_per_doc: int | None = None,
+        instruction: str | None = None,
+    ) -> Mapping[str, object]:
+        if top_n is None and return_documents is None:
+            return CohereRerankParams(
+                query=query,
+                documents=documents,
+            )
+        if top_n is None:
+            return CohereRerankParams(
+                query=query,
+                documents=documents,
+                return_documents=return_documents,
+            )
+        if return_documents is None:
+            return CohereRerankParams(
+                query=query,
+                documents=documents,
+                top_n=top_n,
+            )
+        return CohereRerankParams(
+            query=query,
+            documents=documents,
+            top_n=top_n,
+            return_documents=return_documents,
+        )
+
+    def transform_rerank_request(
+        self,
+        model: str,
+        optional_rerank_params: Mapping[str, object],
+        headers: Mapping[str, object],
+        litellm_params: Mapping[str, object] | None = None,
+    ) -> Mapping[str, object]:
+        query = optional_rerank_params.get("query")
+        documents = optional_rerank_params.get("documents")
+        if not isinstance(query, str) or not query:
+            raise ValueError("query is required for Cloudflare rerank")
+        if not isinstance(documents, Sequence) or isinstance(documents, str) or not documents:
+            raise ValueError("documents is required for Cloudflare rerank")
+
+        contexts = tuple(self._document_to_context(document) for document in documents)
+        request = CloudflareRerankRequest(query=query, contexts=contexts)
+        top_n = optional_rerank_params.get("top_n")
+        if top_n is None:
+            return request
+        return CloudflareRerankRequest(**request, top_k=top_n)
+
+    def transform_rerank_response(
+        self,
+        model: str,
+        raw_response: httpx.Response,
+        model_response: RerankResponse,
+        logging_obj: LiteLLMLoggingObj,
+        api_key: str | None = None,
+        request_data: Mapping[str, object] | None = None,
+        optional_params: Mapping[str, object] | None = None,
+        litellm_params: Mapping[str, object] | None = None,
+    ) -> RerankResponse:
+        request_data = request_data or EMPTY_REQUEST_DATA
+        try:
+            response_json: object = raw_response.json()
+        except ValueError:
+            raise CloudflareError(
+                status_code=raw_response.status_code,
+                message=raw_response.text,
+            )
+        if not isinstance(response_json, Mapping):
+            raise CloudflareError(
+                status_code=raw_response.status_code,
+                message=f"Invalid Cloudflare rerank response: {response_json}",
+            )
+
+        logging_obj.post_call(
+            input=request_data.get("query"),
+            api_key=api_key,
+            additional_args=LoggingAdditionalArgs(complete_input_dict=request_data),
+            original_response=response_json,
+        )
+
+        optional_params = optional_params or EMPTY_REQUEST_DATA
+        documents = optional_params.get("documents")
+        if not isinstance(documents, Sequence) or isinstance(documents, str):
+            documents = ()
+        return_documents = optional_params.get("return_documents") is not False
+        results = tuple(
+            self._transform_response_item(
+                item,
+                documents=documents,
+                return_documents=return_documents,
+            )
+            for item in self._response_items(
+                response_json=response_json,
+                status_code=raw_response.status_code,
+            )
+        )
+        return RerankResponse(
+            id=str(response_json.get("id") or uuid.uuid4()),
+            results=results,
+        )
+
+    def get_error_class(
+        self,
+        error_message: str,
+        status_code: int,
+        headers: Union[Mapping[str, object], httpx.Headers],
+    ) -> BaseLLMException:
+        return CloudflareError(status_code=status_code, message=error_message)

litellm/llms/custom_httpx/llm_http_handler.py

@@ -1209,6 +1209,7 @@ class BaseLLMHTTPHandler:
                 api_key=api_key,
                 timeout=timeout,
                 client=client,
+                optional_rerank_params=optional_rerank_params,
             )
 
         if client is None or not isinstance(client, HTTPHandler):
@@ -1236,6 +1237,7 @@ class BaseLLMHTTPHandler:
             logging_obj=logging_obj,
             api_key=api_key,
             request_data=data,
+            optional_params=optional_rerank_params,
         )
 
     async def arerank(
@@ -1251,6 +1253,7 @@ class BaseLLMHTTPHandler:
         api_key: str | None = None,
         timeout: float | httpx.Timeout | None = None,
         client: HTTPHandler | AsyncHTTPHandler | None = None,
+        optional_rerank_params: Mapping[str, object] | None = None,
     ) -> RerankResponse:
         if client is None or not isinstance(client, AsyncHTTPHandler):
             async_httpx_client = get_async_httpx_client(llm_provider=litellm.LlmProviders(custom_llm_provider))
@@ -1274,6 +1277,7 @@ class BaseLLMHTTPHandler:
             logging_obj=logging_obj,
             api_key=api_key,
             request_data=request_data,
+            optional_params=optional_rerank_params,
         )
 
     def _prepare_audio_transcription_request(

litellm/rerank_api/main.py

@@ -33,7 +33,7 @@ async def arerank(
     query: str,
     documents: list[str | dict[str, Any]],
     custom_llm_provider: (
-        Literal["cohere", "together_ai", "deepinfra", "fireworks_ai", "voyage", "watsonx"] | None
+        Literal["cohere", "together_ai", "deepinfra", "fireworks_ai", "voyage", "watsonx", "cloudflare"] | None
     ) = None,
     top_n: int | None = None,
     rank_fields: list[str] | None = None,
@@ -108,6 +108,7 @@ def rerank(
             "fireworks_ai",
             "voyage",
             "watsonx",
+            "cloudflare",
         ]
         | None
     ) = None,

litellm/utils.py

@@ -8542,6 +8542,8 @@ class ProviderConfigManager:
             )
 
             return get_dashscope_family_rerank_config(provider.value)
+        elif litellm.LlmProviders.CLOUDFLARE == provider:
+            return litellm.CloudflareRerankConfig()
         return litellm.CohereRerankConfig()
 
     @staticmethod

provider_endpoints_support.json

@@ -534,13 +534,13 @@
         "chat_completions": true,
         "messages": true,
         "responses": true,
-        "embeddings": false,
+        "embeddings": true,
         "image_generations": false,
         "audio_transcriptions": false,
         "audio_speech": false,
         "moderations": false,
         "batches": false,
-        "rerank": false,
+        "rerank": true,
         "a2a": true,
         "interactions": true
       }

tests/test_litellm/llms/cloudflare/test_cloudflare_rerank_transformation.py

@@ -0,0 +1,413 @@
+import json
+from unittest.mock import MagicMock, Mock, patch
+
+import httpx
+import pytest
+
+import litellm
+from litellm.llms.cloudflare.chat.transformation import CloudflareError
+from litellm.llms.cloudflare.rerank.transformation import CloudflareRerankConfig
+from litellm.llms.custom_httpx.http_handler import HTTPHandler
+from litellm.types.rerank import RerankResponse
+from litellm.utils import ProviderConfigManager
+
+
+def test_provider_config_manager_returns_cloudflare_rerank_config():
+    config = ProviderConfigManager.get_provider_rerank_config(
+        model="@cf/baai/bge-reranker-base",
+        provider=litellm.LlmProviders.CLOUDFLARE,
+        api_base=None,
+        present_version_params=[],
+    )
+
+    assert isinstance(config, CloudflareRerankConfig)
+
+
+def test_get_complete_url_uses_native_workers_ai_endpoint():
+    config = CloudflareRerankConfig()
+
+    with patch(
+        "litellm.llms.cloudflare.rerank.transformation.get_secret_str",
+        return_value="account-id",
+    ):
+        url = config.get_complete_url(
+            api_base=None,
+            model="@cf/baai/bge-reranker-base",
+        )
+
+    assert url == ("https://api.cloudflare.com/client/v4/accounts/account-id/ai/run/%40cf/baai/bge-reranker-base")
+
+
+def test_get_complete_url_rewrites_openai_compatible_base():
+    config = CloudflareRerankConfig()
+
+    url = config.get_complete_url(
+        api_base="https://api.cloudflare.com/client/v4/accounts/account-id/ai/v1",
+        model="@cf/baai/bge-reranker-base",
+    )
+
+    assert url.endswith("/accounts/account-id/ai/run/%40cf/baai/bge-reranker-base")
+
+
+@pytest.mark.parametrize(
+    ("api_base", "expected"),
+    [
+        (
+            "https://api.cloudflare.com/client/v4/accounts/account-id/ai/run/%40cf/baai/bge-reranker-base",
+            "https://api.cloudflare.com/client/v4/accounts/account-id/ai/run/%40cf/baai/bge-reranker-base",
+        ),
+        (
+            "https://api.cloudflare.com/client/v4/accounts/account-id/ai/run",
+            "https://api.cloudflare.com/client/v4/accounts/account-id/ai/run/%40cf/baai/bge-reranker-base",
+        ),
+        (
+            "https://example.com",
+            "https://example.com/ai/run/%40cf/baai/bge-reranker-base",
+        ),
+    ],
+)
+def test_get_complete_url_handles_supported_base_shapes(api_base, expected):
+    config = CloudflareRerankConfig()
+
+    assert config.get_complete_url(api_base, "@cf/baai/bge-reranker-base") == expected
+
+
+def test_get_complete_url_requires_account_id():
+    config = CloudflareRerankConfig()
+
+    with (
+        patch(
+            "litellm.llms.cloudflare.rerank.transformation.get_secret_str",
+            return_value=None,
+        ),
+        pytest.raises(ValueError, match="CLOUDFLARE_ACCOUNT_ID"),
+    ):
+        config.get_complete_url(None, "@cf/baai/bge-reranker-base")
+
+
+@pytest.mark.parametrize(
+    "model",
+    (
+        "../graphql",
+        "@cf/baai/../graphql",
+        "/@cf/baai/bge-reranker-base",
+    ),
+)
+def test_get_complete_url_rejects_path_traversal(model):
+    config = CloudflareRerankConfig()
+
+    with pytest.raises(ValueError):
+        config.get_complete_url(
+            "https://api.cloudflare.com/client/v4/accounts/account-id/ai/run",
+            model,
+        )
+
+
+def test_get_complete_url_encodes_model_path_segments():
+    config = CloudflareRerankConfig()
+
+    url = config.get_complete_url(
+        "https://api.cloudflare.com/client/v4/accounts/account-id/ai/run",
+        "@cf/baai/model?debug=1",
+    )
+
+    assert url.endswith("/ai/run/%40cf/baai/model%3Fdebug%3D1")
+
+
+def test_validate_environment_and_supported_params():
+    config = CloudflareRerankConfig()
+
+    headers = config.validate_environment(
+        headers={"X-Test": "value"},
+        model="@cf/baai/bge-reranker-base",
+        api_key="cf-key",
+    )
+
+    assert headers["Authorization"] == "Bearer cf-key"
+    assert headers["X-Test"] == "value"
+    assert config.get_supported_cohere_rerank_params("@cf/baai/bge-reranker-base") == (
+        "query",
+        "documents",
+        "top_n",
+        "return_documents",
+    )
+
+
+def test_validate_environment_requires_api_key():
+    config = CloudflareRerankConfig()
+
+    with (
+        patch(
+            "litellm.llms.cloudflare.rerank.transformation.get_secret_str",
+            return_value=None,
+        ),
+        pytest.raises(ValueError, match="Cloudflare API Key"),
+    ):
+        config.validate_environment({}, "@cf/baai/bge-reranker-base")
+
+
+def test_map_cohere_rerank_params_without_top_n():
+    config = CloudflareRerankConfig()
+
+    assert config.map_cohere_rerank_params(
+        non_default_params={},
+        model="@cf/baai/bge-reranker-base",
+        drop_params=False,
+        query="query",
+        documents=("document",),
+    ) == {
+        "query": "query",
+        "documents": ("document",),
+        "return_documents": True,
+    }
+
+
+@pytest.mark.parametrize(
+    ("top_n", "return_documents", "expected"),
+    [
+        (None, None, {"query": "query", "documents": ("document",)}),
+        (
+            1,
+            None,
+            {"query": "query", "documents": ("document",), "top_n": 1},
+        ),
+    ],
+)
+def test_map_cohere_rerank_params_handles_optional_values(top_n, return_documents, expected):
+    config = CloudflareRerankConfig()
+
+    assert (
+        config.map_cohere_rerank_params(
+            non_default_params={},
+            model="@cf/baai/bge-reranker-base",
+            drop_params=False,
+            query="query",
+            documents=("document",),
+            top_n=top_n,
+            return_documents=return_documents,
+        )
+        == expected
+    )
+
+
+def test_transform_rerank_request():
+    config = CloudflareRerankConfig()
+
+    request = config.transform_rerank_request(
+        model="@cf/baai/bge-reranker-base",
+        optional_rerank_params={
+            "query": "Which animal is faster?",
+            "documents": ["A cheetah can sprint.", {"text": "A turtle walks."}],
+            "top_n": 1,
+        },
+        headers={},
+    )
+
+    assert request == {
+        "query": "Which animal is faster?",
+        "contexts": (
+            {"text": "A cheetah can sprint."},
+            {"text": "A turtle walks."},
+        ),
+        "top_k": 1,
+    }
+
+
+@pytest.mark.parametrize(
+    "params",
+    [
+        {"documents": ("document",)},
+        {"query": "query"},
+        {"query": "query", "documents": "document"},
+        {"query": "query", "documents": ()},
+    ],
+)
+def test_transform_rerank_request_validates_required_params(params):
+    config = CloudflareRerankConfig()
+
+    with pytest.raises(ValueError):
+        config.transform_rerank_request(
+            model="@cf/baai/bge-reranker-base",
+            optional_rerank_params=params,
+            headers={},
+        )
+
+
+def test_transform_rerank_request_without_top_n():
+    config = CloudflareRerankConfig()
+
+    request = config.transform_rerank_request(
+        model="@cf/baai/bge-reranker-base",
+        optional_rerank_params={
+            "query": "query",
+            "documents": ({"title": "LiteLLM"},),
+        },
+        headers={},
+    )
+
+    assert request == {
+        "query": "query",
+        "contexts": ({"text": '{"title": "LiteLLM"}'},),
+    }
+
+
+def test_transform_rerank_request_rejects_invalid_document():
+    config = CloudflareRerankConfig()
+
+    with pytest.raises(TypeError, match="strings or dictionaries"):
+        config.transform_rerank_request(
+            model="@cf/baai/bge-reranker-base",
+            optional_rerank_params={
+                "query": "query",
+                "documents": (123,),
+            },
+            headers={},
+        )
+
+
+def test_transform_rerank_response_handles_rest_envelope():
+    config = CloudflareRerankConfig()
+    logging_obj = MagicMock()
+    raw_response = httpx.Response(
+        status_code=200,
+        json={
+            "result": {
+                "response": [
+                    {"id": 1, "score": 0.91},
+                    {"id": 0, "score": 0.42},
+                ]
+            },
+            "success": True,
+            "errors": [],
+            "messages": [],
+        },
+    )
+
+    response = config.transform_rerank_response(
+        model="@cf/baai/bge-reranker-base",
+        raw_response=raw_response,
+        model_response=RerankResponse(),
+        logging_obj=logging_obj,
+        request_data={"query": "query", "contexts": [{"text": "a"}, {"text": "b"}]},
+    )
+
+    assert response.results == [
+        {"index": 1, "relevance_score": 0.91},
+        {"index": 0, "relevance_score": 0.42},
+    ]
+
+
+@pytest.mark.parametrize("return_documents", [True, False])
+def test_transform_rerank_response_honors_return_documents(return_documents):
+    config = CloudflareRerankConfig()
+    raw_response = httpx.Response(
+        status_code=200,
+        json={"result": {"response": [{"id": 0, "score": 0.91}]}},
+    )
+
+    response = config.transform_rerank_response(
+        model="@cf/baai/bge-reranker-base",
+        raw_response=raw_response,
+        model_response=RerankResponse(),
+        logging_obj=MagicMock(),
+        optional_params={
+            "documents": ("LiteLLM is an LLM gateway.",),
+            "return_documents": return_documents,
+        },
+    )
+
+    assert response.results is not None
+    if return_documents:
+        assert response.results[0]["document"] == {"text": "LiteLLM is an LLM gateway."}
+    else:
+        assert "document" not in response.results[0]
+
+
+@pytest.mark.parametrize(
+    ("response_json", "error_type"),
+    [
+        ({"success": False, "errors": ("failed",)}, CloudflareError),
+        ({"result": {}}, CloudflareError),
+        ({"result": {"response": ["invalid"]}}, TypeError),
+        ({"result": {"response": [{"id": "0", "score": 1}]}}, TypeError),
+    ],
+)
+def test_transform_rerank_response_rejects_invalid_responses(response_json, error_type):
+    config = CloudflareRerankConfig()
+    raw_response = httpx.Response(status_code=400, json=response_json)
+
+    with pytest.raises(error_type):
+        config.transform_rerank_response(
+            model="@cf/baai/bge-reranker-base",
+            raw_response=raw_response,
+            model_response=RerankResponse(),
+            logging_obj=MagicMock(),
+        )
+
+
+@pytest.mark.parametrize(
+    "raw_response",
+    [
+        httpx.Response(status_code=500, text="not json"),
+        httpx.Response(status_code=200, json=("not", "an", "object")),
+    ],
+)
+def test_transform_rerank_response_rejects_invalid_json(raw_response):
+    config = CloudflareRerankConfig()
+
+    with pytest.raises(CloudflareError):
+        config.transform_rerank_response(
+            model="@cf/baai/bge-reranker-base",
+            raw_response=raw_response,
+            model_response=RerankResponse(),
+            logging_obj=MagicMock(),
+        )
+
+
+def test_get_error_class():
+    error = CloudflareRerankConfig().get_error_class("failed", 400, {})
+
+    assert isinstance(error, CloudflareError)
+    assert error.status_code == 400
+
+
+def test_litellm_rerank_sends_cloudflare_request():
+    client = HTTPHandler()
+    response_json = {
+        "result": {"response": [{"id": 0, "score": 0.98}]},
+        "success": True,
+    }
+    raw_response = Mock()
+    raw_response.status_code = 200
+    raw_response.json.return_value = response_json
+    raw_response.text = json.dumps(response_json)
+
+    with patch.object(HTTPHandler, "post", return_value=raw_response) as mock_post:
+        response = litellm.rerank(
+            model="cloudflare/@cf/baai/bge-reranker-base",
+            query="What is LiteLLM?",
+            documents=["LiteLLM is an LLM gateway.", "A recipe for soup."],
+            top_n=1,
+            api_key="test-key",
+            api_base="https://api.cloudflare.com/client/v4/accounts/account-id/ai/run",
+            client=client,
+        )
+
+    request = mock_post.call_args.kwargs
+    assert request["url"].endswith("/ai/run/%40cf/baai/bge-reranker-base")
+    assert request["headers"]["Authorization"] == "Bearer test-key"
+    assert json.loads(request["data"]) == {
+        "query": "What is LiteLLM?",
+        "contexts": [
+            {"text": "LiteLLM is an LLM gateway."},
+            {"text": "A recipe for soup."},
+        ],
+        "top_k": 1,
+    }
+    assert response.results == [
+        {
+            "index": 0,
+            "relevance_score": 0.98,
+            "document": {"text": "LiteLLM is an LLM gateway."},
+        }
+    ]