From d8a673092168a417199f9912fef66e6130a9c593 Mon Sep 17 00:00:00 2001 From: prdai Date: Wed, 29 Jul 2026 10:32:16 +0530 Subject: [PATCH 1/9] feat(cloudflare): add embeddings support Co-authored-by: Codex --- README.md | 2 +- litellm/__init__.py | 3 + litellm/_lazy_imports_registry.py | 5 + .../llms/cloudflare/chat/transformation.py | 4 +- .../cloudflare/embedding/transformation.py | 31 ++++ litellm/main.py | 20 +++ litellm/utils.py | 2 + ...est_cloudflare_embedding_transformation.py | 141 ++++++++++++++++++ 8 files changed, 205 insertions(+), 3 deletions(-) create mode 100644 litellm/llms/cloudflare/embedding/transformation.py create mode 100644 tests/test_litellm/llms/cloudflare/test_cloudflare_embedding_transformation.py diff --git a/README.md b/README.md index 31a93bf8c67..e39a622d999 100644 --- a/README.md +++ b/README.md @@ -380,7 +380,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th | [Cerebras (`cerebras`)](https://docs.litellm.ai/docs/providers/cerebras) | ✅ | ✅ | ✅ | | | | | | | | | [Clarifai (`clarifai`)](https://docs.litellm.ai/docs/providers/clarifai) | ✅ | ✅ | ✅ | | | | | | | | | [CLF AI Gateway (`clf_ai_gateway`)](https://docs.litellm.ai/docs/providers/clf_ai_gateway) | ✅ | ✅ | ✅ | | | | | | | | -| [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) | ✅ | ✅ | ✅ | ✅ | | | | | | ✅ | diff --git a/litellm/__init__.py b/litellm/__init__.py index 8e5dd68a186..53d26c424ac 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1841,6 +1841,9 @@ if TYPE_CHECKING: from .llms.cloudflare.chat.transformation import ( CloudflareChatConfig as CloudflareChatConfig, ) + from .llms.cloudflare.embedding.transformation import ( + CloudflareEmbeddingConfig as CloudflareEmbeddingConfig, + ) 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 diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index bf903e23bc0..740bc7f4388 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -192,6 +192,7 @@ LLM_CONFIG_NAMES: Final = ( "LlamaAPIConfig", "TogetherAITextCompletionConfig", "CloudflareChatConfig", + "CloudflareEmbeddingConfig", "NovitaConfig", "PetalsConfig", "OllamaChatConfig", @@ -797,6 +798,10 @@ _LLM_CONFIGS_IMPORT_MAP: Final = { ".llms.cloudflare.chat.transformation", "CloudflareChatConfig", ), + "CloudflareEmbeddingConfig": ( + ".llms.cloudflare.embedding.transformation", + "CloudflareEmbeddingConfig", + ), "NovitaConfig": (".llms.novita.chat.transformation", "NovitaConfig"), "PetalsConfig": (".llms.petals.completion.transformation", "PetalsConfig"), "OllamaChatConfig": (".llms.ollama.chat.transformation", "OllamaChatConfig"), diff --git a/litellm/llms/cloudflare/chat/transformation.py b/litellm/llms/cloudflare/chat/transformation.py index 2bbfd2010cc..51763876bdf 100644 --- a/litellm/llms/cloudflare/chat/transformation.py +++ b/litellm/llms/cloudflare/chat/transformation.py @@ -38,7 +38,7 @@ class CloudflareChatConfig(OpenAIGPTConfig): stream: bool | None = None, ) -> str: return super().get_complete_url( - api_base=self._resolve_api_base(api_base), + api_base=self.resolve_api_base(api_base), api_key=api_key, model=model, optional_params=optional_params, @@ -47,7 +47,7 @@ class CloudflareChatConfig(OpenAIGPTConfig): ) @staticmethod - def _resolve_api_base(api_base: str | None) -> str: + def resolve_api_base(api_base: str | None) -> str: if not api_base: account_id: Final = normalize_nonempty_secret_str(get_secret_str("CLOUDFLARE_ACCOUNT_ID")) if account_id is None: diff --git a/litellm/llms/cloudflare/embedding/transformation.py b/litellm/llms/cloudflare/embedding/transformation.py new file mode 100644 index 00000000000..e3d0c6de911 --- /dev/null +++ b/litellm/llms/cloudflare/embedding/transformation.py @@ -0,0 +1,31 @@ +from collections.abc import Mapping +from typing import Union + +import httpx + +from litellm.llms.cloudflare.chat.transformation import CloudflareChatConfig, CloudflareError +from litellm.llms.vercel_ai_gateway.embedding.transformation import VercelAIGatewayEmbeddingConfig + + +class CloudflareEmbeddingConfig(VercelAIGatewayEmbeddingConfig): + def get_complete_url( + self, + api_base: str | None, + api_key: str | None, + model: str, + optional_params: Mapping[object, object], + litellm_params: Mapping[object, object], + stream: bool | None = None, + ) -> str: + resolved_base = CloudflareChatConfig.resolve_api_base(api_base).rstrip("/") + if resolved_base.endswith("/embeddings"): + return resolved_base + return f"{resolved_base}/embeddings" + + def get_error_class( + self, + error_message: str, + status_code: int, + headers: Union[Mapping[object, object], httpx.Headers], + ) -> CloudflareError: + return CloudflareError(status_code=status_code, message=error_message) diff --git a/litellm/main.py b/litellm/main.py index 6b9964a450f..d6d85ffff60 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7269,6 +7269,26 @@ def embedding( or get_secret_str("VERCEL_OIDC_TOKEN") ) + response = base_llm_http_handler.embedding( + model=model, + input=input, + custom_llm_provider=custom_llm_provider, + api_base=api_base, + api_key=api_key, + logging_obj=logging, + timeout=timeout, + model_response=EmbeddingResponse(), + optional_params=optional_params, + client=client, + aembedding=aembedding, + litellm_params=litellm_params_dict, + headers=headers, + ) + elif custom_llm_provider == "cloudflare": + api_key = api_key or litellm.cloudflare_api_key or litellm.api_key or get_secret_str("CLOUDFLARE_API_KEY") + if api_key is None: + raise ValueError("Missing Cloudflare API Key - no key is set in the environment or request parameters") + api_base = api_base or litellm.api_base or get_secret_str("CLOUDFLARE_API_BASE") response = base_llm_http_handler.embedding( model=model, input=input, diff --git a/litellm/utils.py b/litellm/utils.py index 53b02ebddac..ea6db078b1d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10210,6 +10210,8 @@ class ProviderConfigManager: ) return VercelAIGatewayEmbeddingConfig() + elif litellm.LlmProviders.CLOUDFLARE == provider: + return litellm.CloudflareEmbeddingConfig() elif litellm.LlmProviders.GIGACHAT == provider: return litellm.GigaChatEmbeddingConfig() elif litellm.LlmProviders.HOSTED_VLLM == provider: diff --git a/tests/test_litellm/llms/cloudflare/test_cloudflare_embedding_transformation.py b/tests/test_litellm/llms/cloudflare/test_cloudflare_embedding_transformation.py new file mode 100644 index 00000000000..324df0e5c2a --- /dev/null +++ b/tests/test_litellm/llms/cloudflare/test_cloudflare_embedding_transformation.py @@ -0,0 +1,141 @@ +import json +from unittest.mock import Mock, patch + +import pytest + +import litellm +from litellm.llms.cloudflare.embedding.transformation import CloudflareEmbeddingConfig +from litellm.llms.cloudflare.chat.transformation import CloudflareError +from litellm.llms.custom_httpx.http_handler import HTTPHandler +from litellm.utils import ProviderConfigManager + + +def test_provider_config_manager_returns_cloudflare_embedding_config(): + config = ProviderConfigManager.get_provider_embedding_config( + model="@cf/baai/bge-large-en-v1.5", + provider=litellm.LlmProviders.CLOUDFLARE, + ) + + assert isinstance(config, CloudflareEmbeddingConfig) + + +def test_get_complete_url_defaults_to_openai_compatible_endpoint(monkeypatch): + monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", "acct") + config = CloudflareEmbeddingConfig() + + url = config.get_complete_url( + api_base=None, + api_key="cf-key", + model="@cf/baai/bge-large-en-v1.5", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://api.cloudflare.com/client/v4/accounts/acct/ai/v1/embeddings" + + +def test_get_complete_url_is_idempotent_for_full_endpoint(): + config = CloudflareEmbeddingConfig() + + url = config.get_complete_url( + api_base="https://api.cloudflare.com/client/v4/accounts/acct/ai/v1/embeddings", + api_key="cf-key", + model="@cf/baai/bge-large-en-v1.5", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://api.cloudflare.com/client/v4/accounts/acct/ai/v1/embeddings" + + +def test_get_complete_url_migrates_legacy_ai_run_base(): + config = CloudflareEmbeddingConfig() + + url = config.get_complete_url( + api_base="https://api.cloudflare.com/client/v4/accounts/acct/ai/run/", + api_key="cf-key", + model="@cf/baai/bge-large-en-v1.5", + optional_params={}, + litellm_params={}, + ) + + assert url == "https://api.cloudflare.com/client/v4/accounts/acct/ai/v1/embeddings" + + +def test_validate_environment_preserves_extra_headers(): + config = CloudflareEmbeddingConfig() + + headers = config.validate_environment( + headers={"X-Test": "value"}, + model="@cf/baai/bge-large-en-v1.5", + messages=[], + optional_params={}, + litellm_params={}, + api_key="cf-key", + ) + + assert headers == { + "Authorization": "Bearer cf-key", + "Content-Type": "application/json", + "X-Test": "value", + } + + +def test_get_error_class(): + error = CloudflareEmbeddingConfig().get_error_class("failed", 400, {}) + + assert isinstance(error, CloudflareError) + assert error.status_code == 400 + + +def test_embedding_routes_to_cloudflare_openai_compatible_endpoint(monkeypatch): + monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", "acct") + client = HTTPHandler() + response_json = { + "object": "list", + "data": [ + { + "object": "embedding", + "embedding": [0.1, 0.2, 0.3], + "index": 0, + } + ], + "model": "@cf/baai/bge-large-en-v1.5", + "usage": {"prompt_tokens": 2, "total_tokens": 2}, + } + raw_response = Mock() + raw_response.status_code = 200 + raw_response.headers = {"content-type": "application/json"} + 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.embedding( + model="cloudflare/@cf/baai/bge-large-en-v1.5", + input=["hello"], + api_key="cf-key", + client=client, + caching=False, + ) + + request = mock_post.call_args.kwargs + body = json.loads(request["data"]) + assert request["url"] == "https://api.cloudflare.com/client/v4/accounts/acct/ai/v1/embeddings" + assert request["headers"]["Authorization"] == "Bearer cf-key" + assert body == { + "model": "@cf/baai/bge-large-en-v1.5", + "input": ["hello"], + } + assert response.data[0]["embedding"] == [0.1, 0.2, 0.3] + + +def test_embedding_requires_cloudflare_api_key(monkeypatch): + monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", "acct") + monkeypatch.delenv("CLOUDFLARE_API_KEY", raising=False) + + with pytest.raises(litellm.APIConnectionError, match="Missing Cloudflare API Key"): + litellm.embedding( + model="cloudflare/@cf/baai/bge-large-en-v1.5", + input=["hello"], + caching=False, + ) From 949c9be3ac6b18723cdd29c9676cc4821c25b42c Mon Sep 17 00:00:00 2001 From: prdai Date: Wed, 29 Jul 2026 10:53:59 +0530 Subject: [PATCH 2/9] feat(cloudflare): add rerank support Co-authored-by: Codex --- README.md | 2 +- litellm/__init__.py | 3 + litellm/_lazy_imports_registry.py | 5 + .../llms/cloudflare/rerank/transformation.py | 284 ++++++++++++ litellm/llms/custom_httpx/llm_http_handler.py | 4 + litellm/rerank_api/main.py | 3 +- litellm/utils.py | 2 + provider_endpoints_support.json | 4 +- .../test_cloudflare_rerank_transformation.py | 413 ++++++++++++++++++ 9 files changed, 716 insertions(+), 4 deletions(-) create mode 100644 litellm/llms/cloudflare/rerank/transformation.py create mode 100644 tests/test_litellm/llms/cloudflare/test_cloudflare_rerank_transformation.py diff --git a/README.md b/README.md index e39a622d999..b1de5c7033a 100644 --- a/README.md +++ b/README.md @@ -380,7 +380,7 @@ Set `LITELLM_PROXY_API_BASE` and `LITELLM_PROXY_API_KEY` and every model call th | [Cerebras (`cerebras`)](https://docs.litellm.ai/docs/providers/cerebras) | ✅ | ✅ | ✅ | | | | | | | | | [Clarifai (`clarifai`)](https://docs.litellm.ai/docs/providers/clarifai) | ✅ | ✅ | ✅ | | | | | | | | | [CLF AI Gateway (`clf_ai_gateway`)](https://docs.litellm.ai/docs/providers/clf_ai_gateway) | ✅ | ✅ | ✅ | | | | | | | | -| [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) | ✅ | ✅ | ✅ | ✅ | | | | | | ✅ | diff --git a/litellm/__init__.py b/litellm/__init__.py index 53d26c424ac..711a7f8ecc7 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1844,6 +1844,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 diff --git a/litellm/_lazy_imports_registry.py b/litellm/_lazy_imports_registry.py index 740bc7f4388..1ef819c76c3 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -193,6 +193,7 @@ LLM_CONFIG_NAMES: Final = ( "TogetherAITextCompletionConfig", "CloudflareChatConfig", "CloudflareEmbeddingConfig", + "CloudflareRerankConfig", "NovitaConfig", "PetalsConfig", "OllamaChatConfig", @@ -802,6 +803,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"), diff --git a/litellm/llms/cloudflare/rerank/transformation.py b/litellm/llms/cloudflare/rerank/transformation.py new file mode 100644 index 00000000000..b26f7843513 --- /dev/null +++ b/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) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 4a10db31606..8e7cf857be3 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1611,6 +1611,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): @@ -1638,6 +1639,7 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, api_key=api_key, request_data=data, + optional_params=optional_rerank_params, ) async def arerank( @@ -1653,6 +1655,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)) @@ -1676,6 +1679,7 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, api_key=api_key, request_data=request_data, + optional_params=optional_rerank_params, ) def _prepare_decisions_request( diff --git a/litellm/rerank_api/main.py b/litellm/rerank_api/main.py index e344ce46be7..8e224f5db1b 100644 --- a/litellm/rerank_api/main.py +++ b/litellm/rerank_api/main.py @@ -39,7 +39,7 @@ async def arerank( query: str, documents: list[str | dict[str, object]], 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, @@ -112,6 +112,7 @@ def rerank( "fireworks_ai", "voyage", "watsonx", + "cloudflare", ] | None ) = None, diff --git a/litellm/utils.py b/litellm/utils.py index ea6db078b1d..b4c7962cfe6 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10280,6 +10280,8 @@ class ProviderConfigManager: ) return get_dashscope_family_rerank_config(provider.value) + elif litellm.LlmProviders.CLOUDFLARE == provider: + return litellm.CloudflareRerankConfig() return litellm.CohereRerankConfig() @staticmethod diff --git a/provider_endpoints_support.json b/provider_endpoints_support.json index 00fee0a0d5a..9fe38cf68e1 100644 --- a/provider_endpoints_support.json +++ b/provider_endpoints_support.json @@ -860,12 +860,12 @@ "chat_completions": true, "responses": true, "text_completion": true, - "embeddings": false, + "embeddings": true, "image_generations": false, "image_edits": false, "audio_transcriptions": true, "audio_speech": false, - "rerank": false, + "rerank": true, "ocr": false, "search": false, "moderations": false, diff --git a/tests/test_litellm/llms/cloudflare/test_cloudflare_rerank_transformation.py b/tests/test_litellm/llms/cloudflare/test_cloudflare_rerank_transformation.py new file mode 100644 index 00000000000..d4b1c681f2e --- /dev/null +++ b/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."}, + } + ] From ff14a322a4a92d2453e2cc069c3e1ab8349828e3 Mon Sep 17 00:00:00 2001 From: prdai Date: Sun, 13 Sep 2026 15:05:26 +0530 Subject: [PATCH 3/9] fix(cloudflare): accept litellm_params in rerank validate_environment --- litellm/llms/cloudflare/rerank/transformation.py | 1 + 1 file changed, 1 insertion(+) diff --git a/litellm/llms/cloudflare/rerank/transformation.py b/litellm/llms/cloudflare/rerank/transformation.py index b26f7843513..e8f776b05cd 100644 --- a/litellm/llms/cloudflare/rerank/transformation.py +++ b/litellm/llms/cloudflare/rerank/transformation.py @@ -118,6 +118,7 @@ class CloudflareRerankConfig(BaseRerankConfig): model: str, api_key: str | None = None, optional_params: Mapping[str, object] | None = None, + litellm_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: From 23c9fc2e01f9604e45bcdd65934db50461f3a6ec Mon Sep 17 00:00:00 2001 From: prdai Date: Sun, 13 Sep 2026 15:11:51 +0530 Subject: [PATCH 4/9] style(cloudflare): apply ruff union and formatting fixes --- litellm/llms/cloudflare/embedding/transformation.py | 3 +-- litellm/llms/cloudflare/rerank/transformation.py | 12 +++--------- 2 files changed, 4 insertions(+), 11 deletions(-) diff --git a/litellm/llms/cloudflare/embedding/transformation.py b/litellm/llms/cloudflare/embedding/transformation.py index e3d0c6de911..aa62526322e 100644 --- a/litellm/llms/cloudflare/embedding/transformation.py +++ b/litellm/llms/cloudflare/embedding/transformation.py @@ -1,5 +1,4 @@ from collections.abc import Mapping -from typing import Union import httpx @@ -26,6 +25,6 @@ class CloudflareEmbeddingConfig(VercelAIGatewayEmbeddingConfig): self, error_message: str, status_code: int, - headers: Union[Mapping[object, object], httpx.Headers], + headers: Mapping[object, object] | httpx.Headers, ) -> CloudflareError: return CloudflareError(status_code=status_code, message=error_message) diff --git a/litellm/llms/cloudflare/rerank/transformation.py b/litellm/llms/cloudflare/rerank/transformation.py index e8f776b05cd..5ff3f2dca38 100644 --- a/litellm/llms/cloudflare/rerank/transformation.py +++ b/litellm/llms/cloudflare/rerank/transformation.py @@ -1,6 +1,5 @@ import json from collections.abc import Mapping, Sequence -from typing import Union import httpx from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -28,7 +27,7 @@ class CloudflareRerankRequest(TypedDict): class CohereRerankParams(TypedDict): query: ReadOnly[str] - documents: ReadOnly[Sequence[Union[str, Mapping[str, object]]]] + documents: ReadOnly[Sequence[str | Mapping[str, object]]] top_n: NotRequired[ReadOnly[int]] return_documents: NotRequired[ReadOnly[bool]] @@ -164,12 +163,7 @@ class CloudflareRerankConfig(BaseRerankConfig): model: str, drop_params: bool, query: str, - documents: Sequence[ - Union[ - str, - Mapping[str, object], - ] - ], + documents: Sequence[str | Mapping[str, object]], custom_llm_provider: str | None = None, top_n: int | None = None, rank_fields: Sequence[str] | None = None, @@ -280,6 +274,6 @@ class CloudflareRerankConfig(BaseRerankConfig): self, error_message: str, status_code: int, - headers: Union[Mapping[str, object], httpx.Headers], + headers: Mapping[str, object] | httpx.Headers, ) -> BaseLLMException: return CloudflareError(status_code=status_code, message=error_message) From badd5c8e9a2d7dc05db3a4c7950feafd56878260 Mon Sep 17 00:00:00 2001 From: prdai Date: Sun, 13 Sep 2026 15:24:48 +0530 Subject: [PATCH 5/9] chore(cloudflare): satisfy type-discipline and complexity gates --- .../llms/cloudflare/rerank/transformation.py | 19 +++++++------------ litellm/utils.py | 2 +- 2 files changed, 8 insertions(+), 13 deletions(-) diff --git a/litellm/llms/cloudflare/rerank/transformation.py b/litellm/llms/cloudflare/rerank/transformation.py index 5ff3f2dca38..17a5a840bf3 100644 --- a/litellm/llms/cloudflare/rerank/transformation.py +++ b/litellm/llms/cloudflare/rerank/transformation.py @@ -1,5 +1,6 @@ import json from collections.abc import Mapping, Sequence +from typing import Final import httpx from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -32,11 +33,6 @@ class CohereRerankParams(TypedDict): return_documents: NotRequired[ReadOnly[bool]] -class CloudflareHeaders(TypedDict): - Authorization: str - accept: str - - class LoggingAdditionalArgs(TypedDict): complete_input_dict: ReadOnly[Mapping[str, object]] @@ -122,13 +118,12 @@ class CloudflareRerankConfig(BaseRerankConfig): 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 + default_headers: Final = { + "Authorization": f"Bearer {api_key}", + "accept": "application/json", + "content-type": "application/json", + } + return {**default_headers, **headers} def get_complete_url( self, diff --git a/litellm/utils.py b/litellm/utils.py index b4c7962cfe6..84d7c0d151d 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -10229,7 +10229,7 @@ class ProviderConfigManager: return None @staticmethod - def get_provider_rerank_config( + def get_provider_rerank_config( # noqa: C901 # provider dispatch; one branch per rerank provider model: str, provider: LlmProviders, api_base: str | None, From c6a917827d8be36b87d441b8df80154721bcd845 Mon Sep 17 00:00:00 2001 From: prdai Date: Sun, 13 Sep 2026 15:24:48 +0530 Subject: [PATCH 6/9] test(cloudflare): fake the HTTP boundary instead of patching litellm internals --- ...est_cloudflare_embedding_transformation.py | 33 ++++-- .../test_cloudflare_rerank_transformation.py | 107 +++++++++--------- 2 files changed, 76 insertions(+), 64 deletions(-) diff --git a/tests/test_litellm/llms/cloudflare/test_cloudflare_embedding_transformation.py b/tests/test_litellm/llms/cloudflare/test_cloudflare_embedding_transformation.py index 324df0e5c2a..17ce0f83a18 100644 --- a/tests/test_litellm/llms/cloudflare/test_cloudflare_embedding_transformation.py +++ b/tests/test_litellm/llms/cloudflare/test_cloudflare_embedding_transformation.py @@ -1,5 +1,5 @@ import json -from unittest.mock import Mock, patch +from unittest.mock import Mock import pytest @@ -10,6 +10,18 @@ from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.utils import ProviderConfigManager +class _RecordingHTTPHandler(HTTPHandler): + def __init__(self, response): + super().__init__() + self.response = response + self.requests = [] + + def post(self, url: str, **kwargs): + self.requests.append({"url": url, **kwargs}) + return self.response + + + def test_provider_config_manager_returns_cloudflare_embedding_config(): config = ProviderConfigManager.get_provider_embedding_config( model="@cf/baai/bge-large-en-v1.5", @@ -90,7 +102,6 @@ def test_get_error_class(): def test_embedding_routes_to_cloudflare_openai_compatible_endpoint(monkeypatch): monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", "acct") - client = HTTPHandler() response_json = { "object": "list", "data": [ @@ -108,17 +119,17 @@ def test_embedding_routes_to_cloudflare_openai_compatible_endpoint(monkeypatch): raw_response.headers = {"content-type": "application/json"} raw_response.json.return_value = response_json raw_response.text = json.dumps(response_json) + client = _RecordingHTTPHandler(raw_response) - with patch.object(HTTPHandler, "post", return_value=raw_response) as mock_post: - response = litellm.embedding( - model="cloudflare/@cf/baai/bge-large-en-v1.5", - input=["hello"], - api_key="cf-key", - client=client, - caching=False, - ) + response = litellm.embedding( + model="cloudflare/@cf/baai/bge-large-en-v1.5", + input=["hello"], + api_key="cf-key", + client=client, + caching=False, + ) - request = mock_post.call_args.kwargs + request = client.requests[0] body = json.loads(request["data"]) assert request["url"] == "https://api.cloudflare.com/client/v4/accounts/acct/ai/v1/embeddings" assert request["headers"]["Authorization"] == "Bearer cf-key" diff --git a/tests/test_litellm/llms/cloudflare/test_cloudflare_rerank_transformation.py b/tests/test_litellm/llms/cloudflare/test_cloudflare_rerank_transformation.py index d4b1c681f2e..ff95603bf2a 100644 --- a/tests/test_litellm/llms/cloudflare/test_cloudflare_rerank_transformation.py +++ b/tests/test_litellm/llms/cloudflare/test_cloudflare_rerank_transformation.py @@ -1,5 +1,5 @@ import json -from unittest.mock import MagicMock, Mock, patch +from unittest.mock import MagicMock, Mock import httpx import pytest @@ -12,6 +12,18 @@ from litellm.types.rerank import RerankResponse from litellm.utils import ProviderConfigManager +class _RecordingHTTPHandler(HTTPHandler): + def __init__(self, response): + super().__init__() + self.response = response + self.requests = [] + + def post(self, url: str, **kwargs): + self.requests.append({"url": url, **kwargs}) + return self.response + + + def test_provider_config_manager_returns_cloudflare_rerank_config(): config = ProviderConfigManager.get_provider_rerank_config( model="@cf/baai/bge-reranker-base", @@ -23,17 +35,14 @@ def test_provider_config_manager_returns_cloudflare_rerank_config(): assert isinstance(config, CloudflareRerankConfig) -def test_get_complete_url_uses_native_workers_ai_endpoint(): +def test_get_complete_url_uses_native_workers_ai_endpoint(monkeypatch): + monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", "account-id") 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", - ) + 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") @@ -72,31 +81,26 @@ def test_get_complete_url_handles_supported_base_shapes(api_base, expected): assert config.get_complete_url(api_base, "@cf/baai/bge-reranker-base") == expected -def test_get_complete_url_requires_account_id(): +def test_get_complete_url_requires_account_id(monkeypatch): + monkeypatch.delenv("CLOUDFLARE_ACCOUNT_ID", raising=False) config = CloudflareRerankConfig() - with ( - patch( - "litellm.llms.cloudflare.rerank.transformation.get_secret_str", - return_value=None, - ), - pytest.raises(ValueError, match="CLOUDFLARE_ACCOUNT_ID"), - ): + with pytest.raises(ValueError, match="CLOUDFLARE_ACCOUNT_ID"): config.get_complete_url(None, "@cf/baai/bge-reranker-base") @pytest.mark.parametrize( - "model", + "model,error_match", ( - "../graphql", - "@cf/baai/../graphql", - "/@cf/baai/bge-reranker-base", + ("../graphql", "cannot be a dot path segment"), + ("@cf/baai/../graphql", "cannot be a dot path segment"), + ("/@cf/baai/bge-reranker-base", "model is required"), ), ) -def test_get_complete_url_rejects_path_traversal(model): +def test_get_complete_url_rejects_path_traversal(model, error_match): config = CloudflareRerankConfig() - with pytest.raises(ValueError): + with pytest.raises(ValueError, match=error_match): config.get_complete_url( "https://api.cloudflare.com/client/v4/accounts/account-id/ai/run", model, @@ -133,16 +137,11 @@ def test_validate_environment_and_supported_params(): ) -def test_validate_environment_requires_api_key(): +def test_validate_environment_requires_api_key(monkeypatch): + monkeypatch.delenv("CLOUDFLARE_API_KEY", raising=False) config = CloudflareRerankConfig() - with ( - patch( - "litellm.llms.cloudflare.rerank.transformation.get_secret_str", - return_value=None, - ), - pytest.raises(ValueError, match="Cloudflare API Key"), - ): + with pytest.raises(ValueError, match="Cloudflare API Key"): config.validate_environment({}, "@cf/baai/bge-reranker-base") @@ -214,18 +213,21 @@ def test_transform_rerank_request(): @pytest.mark.parametrize( - "params", - [ - {"documents": ("document",)}, - {"query": "query"}, - {"query": "query", "documents": "document"}, - {"query": "query", "documents": ()}, - ], + "params,error_match", + ( + ({"documents": ("document",)}, "query is required for Cloudflare rerank"), + ({"query": "query"}, "documents is required for Cloudflare rerank"), + ( + {"query": "query", "documents": "document"}, + "documents is required for Cloudflare rerank", + ), + ({"query": "query", "documents": ()}, "documents is required for Cloudflare rerank"), + ), ) -def test_transform_rerank_request_validates_required_params(params): +def test_transform_rerank_request_validates_required_params(params, error_match): config = CloudflareRerankConfig() - with pytest.raises(ValueError): + with pytest.raises(ValueError, match=error_match): config.transform_rerank_request( model="@cf/baai/bge-reranker-base", optional_rerank_params=params, @@ -372,7 +374,6 @@ def test_get_error_class(): def test_litellm_rerank_sends_cloudflare_request(): - client = HTTPHandler() response_json = { "result": {"response": [{"id": 0, "score": 0.98}]}, "success": True, @@ -381,19 +382,19 @@ def test_litellm_rerank_sends_cloudflare_request(): raw_response.status_code = 200 raw_response.json.return_value = response_json raw_response.text = json.dumps(response_json) + client = _RecordingHTTPHandler(raw_response) - 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, - ) + 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 + request = client.requests[0] assert request["url"].endswith("/ai/run/%40cf/baai/bge-reranker-base") assert request["headers"]["Authorization"] == "Bearer test-key" assert json.loads(request["data"]) == { From bc21e9dd4324a8b10ebdfc2b0eb9ddea2a945aa2 Mon Sep 17 00:00:00 2001 From: prdai Date: Sun, 13 Sep 2026 15:45:25 +0530 Subject: [PATCH 7/9] fix(cloudflare): satisfy basedpyright override and complexity gates --- .../cloudflare/embedding/transformation.py | 5 +- .../llms/cloudflare/rerank/transformation.py | 106 ++++++++++-------- litellm/llms/custom_httpx/llm_http_handler.py | 7 +- litellm/main.py | 62 +++++----- 4 files changed, 95 insertions(+), 85 deletions(-) diff --git a/litellm/llms/cloudflare/embedding/transformation.py b/litellm/llms/cloudflare/embedding/transformation.py index aa62526322e..2d75c8dab98 100644 --- a/litellm/llms/cloudflare/embedding/transformation.py +++ b/litellm/llms/cloudflare/embedding/transformation.py @@ -1,4 +1,5 @@ from collections.abc import Mapping +from typing import Final import httpx @@ -16,7 +17,7 @@ class CloudflareEmbeddingConfig(VercelAIGatewayEmbeddingConfig): litellm_params: Mapping[object, object], stream: bool | None = None, ) -> str: - resolved_base = CloudflareChatConfig.resolve_api_base(api_base).rstrip("/") + resolved_base: Final = CloudflareChatConfig.resolve_api_base(api_base).rstrip("/") if resolved_base.endswith("/embeddings"): return resolved_base return f"{resolved_base}/embeddings" @@ -25,6 +26,6 @@ class CloudflareEmbeddingConfig(VercelAIGatewayEmbeddingConfig): self, error_message: str, status_code: int, - headers: Mapping[object, object] | httpx.Headers, + headers: Mapping[str, str] | httpx.Headers, ) -> CloudflareError: return CloudflareError(status_code=status_code, message=error_message) diff --git a/litellm/llms/cloudflare/rerank/transformation.py b/litellm/llms/cloudflare/rerank/transformation.py index 17a5a840bf3..4cc9b4f374a 100644 --- a/litellm/llms/cloudflare/rerank/transformation.py +++ b/litellm/llms/cloudflare/rerank/transformation.py @@ -1,6 +1,6 @@ import json from collections.abc import Mapping, Sequence -from typing import Final +from typing import Final, cast # noqa: TID251 # json boundaries only expose untyped mappings import httpx from typing_extensions import NotRequired, ReadOnly, TypedDict @@ -11,7 +11,7 @@ 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 litellm.types.rerank import RerankResponse, RerankResponseDocument, RerankResponseResult from ..chat.transformation import CloudflareError @@ -22,7 +22,7 @@ class CloudflareRerankContext(TypedDict): class CloudflareRerankRequest(TypedDict): query: ReadOnly[str] - contexts: ReadOnly[Sequence[Mapping[str, str]]] + contexts: ReadOnly[Sequence[CloudflareRerankContext]] top_k: NotRequired[ReadOnly[object]] @@ -41,13 +41,13 @@ class EmptyRequestData(TypedDict, total=False): pass -EMPTY_REQUEST_DATA: Mapping[str, object] = EmptyRequestData() +EMPTY_REQUEST_DATA: Final[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")) + account_id: Final = 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" @@ -55,14 +55,15 @@ class CloudflareRerankConfig(BaseRerankConfig): @staticmethod def _document_to_context( document: object, - ) -> Mapping[str, str]: + ) -> CloudflareRerankContext: 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)) + document_mapping: Final = cast(Mapping[str, object], document) # cast-ok: isinstance proves only a bare Mapping + text: Final = document_mapping.get("text") + return CloudflareRerankContext(text=text if isinstance(text, str) else json.dumps(document_mapping)) @staticmethod def _response_items( @@ -75,14 +76,19 @@ class CloudflareRerankConfig(BaseRerankConfig): 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 + result: Final = response_json.get("result", response_json) + result_mapping: Final = ( + cast(Mapping[str, object], result) # cast-ok: isinstance proves only a bare Mapping + if isinstance(result, Mapping) + else None + ) + response: Final = result_mapping.get("response") if result_mapping is not None 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) + return tuple(cast(Sequence[object], response)) # cast-ok: JSON list narrows only to list[Unknown] @staticmethod def _transform_response_item( @@ -93,21 +99,22 @@ class CloudflareRerankConfig(BaseRerankConfig): if not isinstance(item, Mapping): raise TypeError("Invalid item in Cloudflare rerank response") - index = item.get("id") - score = item.get("score") + item_mapping: Final = cast(Mapping[str, object], item) # cast-ok: isinstance proves only a bare Mapping + index: Final = item_mapping.get("id") + score: Final = item_mapping.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]) + document: Final = CloudflareRerankConfig._document_to_context(documents[index]) return RerankResponseResult( index=index, relevance_score=float(score), - document=document, + document=RerankResponseDocument(text=document["text"]), ) return RerankResponseResult(index=index, relevance_score=float(score)) - def validate_environment( + def validate_environment( # pyright: ignore[reportIncompatibleMethodOverride] # base annotates the return as dict; the read-only Mapping is intentional self, headers: Mapping[str, object], model: str, @@ -115,11 +122,11 @@ class CloudflareRerankConfig(BaseRerankConfig): optional_params: Mapping[str, object] | None = None, litellm_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: + resolved_api_key: Final = api_key or get_secret_str("CLOUDFLARE_API_KEY") + if resolved_api_key is None: raise ValueError("Missing Cloudflare API Key - set CLOUDFLARE_API_KEY or pass api_key explicitly") default_headers: Final = { - "Authorization": f"Bearer {api_key}", + "Authorization": f"Bearer {resolved_api_key}", "accept": "application/json", "content-type": "application/json", } @@ -131,8 +138,8 @@ class CloudflareRerankConfig(BaseRerankConfig): 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") + cleaned: Final = (api_base or self._default_api_base()).rstrip("/") + encoded_model: Final = encode_url_path_segments(model, field_name="model") if cleaned.endswith(f"/{encoded_model}"): return cleaned if cleaned.endswith("/ai/v1"): @@ -141,7 +148,7 @@ class CloudflareRerankConfig(BaseRerankConfig): return f"{cleaned}/{encoded_model}" return f"{cleaned}/ai/run/{encoded_model}" - def get_supported_cohere_rerank_params( + def get_supported_cohere_rerank_params( # pyright: ignore[reportIncompatibleMethodOverride] # base annotates the return as list; a read-only tuple is intentional self, model: str, ) -> Sequence[str]: @@ -152,7 +159,7 @@ class CloudflareRerankConfig(BaseRerankConfig): "return_documents", ) - def map_cohere_rerank_params( + def map_cohere_rerank_params( # pyright: ignore[reportIncompatibleMethodOverride] # base annotates the return as dict; the read-only Mapping is intentional self, non_default_params: Mapping[str, object], model: str, @@ -173,6 +180,7 @@ class CloudflareRerankConfig(BaseRerankConfig): documents=documents, ) if top_n is None: + assert return_documents is not None return CohereRerankParams( query=query, documents=documents, @@ -191,26 +199,26 @@ class CloudflareRerankConfig(BaseRerankConfig): return_documents=return_documents, ) - def transform_rerank_request( + def transform_rerank_request( # pyright: ignore[reportIncompatibleMethodOverride] # base annotates the return as dict; the read-only Mapping is intentional 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") + query: Final = optional_rerank_params.get("query") + raw_documents: Final = 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: + if not isinstance(raw_documents, Sequence) or isinstance(raw_documents, str) or not raw_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") + documents: Final = cast(Sequence[object], raw_documents) # cast-ok: isinstance proves only a bare Sequence + contexts: Final = tuple(self._document_to_context(document) for document in documents) + top_n: Final = optional_rerank_params.get("top_n") if top_n is None: - return request - return CloudflareRerankRequest(**request, top_k=top_n) + return CloudflareRerankRequest(query=query, contexts=contexts) + return CloudflareRerankRequest(query=query, contexts=contexts, top_k=top_n) def transform_rerank_response( self, @@ -223,9 +231,9 @@ class CloudflareRerankConfig(BaseRerankConfig): optional_params: Mapping[str, object] | None = None, litellm_params: Mapping[str, object] | None = None, ) -> RerankResponse: - request_data = request_data or EMPTY_REQUEST_DATA + resolved_request_data: Final = request_data or EMPTY_REQUEST_DATA try: - response_json: object = raw_response.json() + response_json: Final = cast(object, raw_response.json()) # cast-ok: httpx Response.json is untyped except ValueError: raise CloudflareError( status_code=raw_response.status_code, @@ -236,33 +244,37 @@ class CloudflareRerankConfig(BaseRerankConfig): status_code=raw_response.status_code, message=f"Invalid Cloudflare rerank response: {response_json}", ) + response_data: Final = cast(Mapping[str, object], response_json) # cast-ok: bare Mapping - logging_obj.post_call( - input=request_data.get("query"), + logging_obj.post_call( # pyright: ignore[reportUnknownMemberType] # logging object is legacy-untyped + input=resolved_request_data.get("query"), api_key=api_key, - additional_args=LoggingAdditionalArgs(complete_input_dict=request_data), - original_response=response_json, + additional_args=LoggingAdditionalArgs(complete_input_dict=resolved_request_data), + original_response=response_data, ) - 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( + resolved_optional_params: Final = optional_params or EMPTY_REQUEST_DATA + raw_documents: Final = resolved_optional_params.get("documents") + documents: Final = ( + cast(Sequence[object], raw_documents) # cast-ok: isinstance proves only a bare Sequence + if isinstance(raw_documents, Sequence) and not isinstance(raw_documents, str) + else () + ) + return_documents: Final = resolved_optional_params.get("return_documents") is not False + results: Final = tuple( self._transform_response_item( item, documents=documents, return_documents=return_documents, ) for item in self._response_items( - response_json=response_json, + response_json=response_data, status_code=raw_response.status_code, ) ) return RerankResponse( - id=str(response_json.get("id") or uuid.uuid4()), - results=results, + id=str(response_data.get("id") or uuid.uuid4()), + results=list(results), ) def get_error_class( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 8e7cf857be3..880d5322a1e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -1611,7 +1611,9 @@ class BaseLLMHTTPHandler: api_key=api_key, timeout=timeout, client=client, - optional_rerank_params=optional_rerank_params, + optional_rerank_params=cast( # cast-ok: untyped caller dict + Mapping[str, object], optional_rerank_params + ), ) if client is None or not isinstance(client, HTTPHandler): @@ -1672,6 +1674,7 @@ class BaseLLMHTTPHandler: except Exception as e: raise self._handle_error(e=e, provider_config=provider_config) + resolved_optional_rerank_params: Final = dict(optional_rerank_params or {}) return provider_config.transform_rerank_response( model=model, raw_response=response, @@ -1679,7 +1682,7 @@ class BaseLLMHTTPHandler: logging_obj=logging_obj, api_key=api_key, request_data=request_data, - optional_params=optional_rerank_params, + optional_params=resolved_optional_rerank_params, ) def _prepare_decisions_request( diff --git a/litellm/main.py b/litellm/main.py index d6d85ffff60..05941ea57ba 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6793,6 +6793,31 @@ def _embedding_headers(param: str, configured: object, model: str, drop_params: ) + +def _resolve_vercel_or_cloudflare_embedding_credentials( + custom_llm_provider: str, + api_base: str | None, + api_key: str | None, +) -> tuple[str | None, str | None]: + if custom_llm_provider == "cloudflare": + resolved_api_key: Final = ( + api_key or litellm.cloudflare_api_key or litellm.api_key or get_secret_str("CLOUDFLARE_API_KEY") + ) + if resolved_api_key is None: + raise ValueError("Missing Cloudflare API Key - no key is set in the environment or request parameters") + return api_base or litellm.api_base or get_secret_str("CLOUDFLARE_API_BASE"), resolved_api_key + return ( + api_base + or litellm.api_base + or get_secret_str("VERCEL_AI_GATEWAY_API_BASE") + or "https://ai-gateway.vercel.sh/v1", + api_key + or litellm.api_key + or get_secret_str("VERCEL_AI_GATEWAY_API_KEY") + or get_secret_str("VERCEL_OIDC_TOKEN"), + ) + + # fmt: off # Overload for when aembedding=True (returns coroutine) @@ -7254,41 +7279,10 @@ def embedding( litellm_params=litellm_params_dict, headers=headers, ) - elif custom_llm_provider == "vercel_ai_gateway": - api_base = ( - api_base - or litellm.api_base - or get_secret_str("VERCEL_AI_GATEWAY_API_BASE") - or "https://ai-gateway.vercel.sh/v1" + elif custom_llm_provider in ("vercel_ai_gateway", "cloudflare"): + api_base, api_key = _resolve_vercel_or_cloudflare_embedding_credentials( + custom_llm_provider, api_base, api_key ) - - api_key = ( - api_key - or litellm.api_key - or get_secret_str("VERCEL_AI_GATEWAY_API_KEY") - or get_secret_str("VERCEL_OIDC_TOKEN") - ) - - response = base_llm_http_handler.embedding( - model=model, - input=input, - custom_llm_provider=custom_llm_provider, - api_base=api_base, - api_key=api_key, - logging_obj=logging, - timeout=timeout, - model_response=EmbeddingResponse(), - optional_params=optional_params, - client=client, - aembedding=aembedding, - litellm_params=litellm_params_dict, - headers=headers, - ) - elif custom_llm_provider == "cloudflare": - api_key = api_key or litellm.cloudflare_api_key or litellm.api_key or get_secret_str("CLOUDFLARE_API_KEY") - if api_key is None: - raise ValueError("Missing Cloudflare API Key - no key is set in the environment or request parameters") - api_base = api_base or litellm.api_base or get_secret_str("CLOUDFLARE_API_BASE") response = base_llm_http_handler.embedding( model=model, input=input, From bc796d7971408c4c1677aefecbace2da8e99fd58 Mon Sep 17 00:00:00 2001 From: prdai Date: Fri, 6 Nov 2026 13:25:46 +0530 Subject: [PATCH 8/9] test(cloudflare): move provider tests under tests/unit --- .../llms/cloudflare/test_cloudflare_embedding_transformation.py | 0 .../llms/cloudflare/test_cloudflare_rerank_transformation.py | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename tests/{test_litellm => unit}/llms/cloudflare/test_cloudflare_embedding_transformation.py (100%) rename tests/{test_litellm => unit}/llms/cloudflare/test_cloudflare_rerank_transformation.py (100%) diff --git a/tests/test_litellm/llms/cloudflare/test_cloudflare_embedding_transformation.py b/tests/unit/llms/cloudflare/test_cloudflare_embedding_transformation.py similarity index 100% rename from tests/test_litellm/llms/cloudflare/test_cloudflare_embedding_transformation.py rename to tests/unit/llms/cloudflare/test_cloudflare_embedding_transformation.py diff --git a/tests/test_litellm/llms/cloudflare/test_cloudflare_rerank_transformation.py b/tests/unit/llms/cloudflare/test_cloudflare_rerank_transformation.py similarity index 100% rename from tests/test_litellm/llms/cloudflare/test_cloudflare_rerank_transformation.py rename to tests/unit/llms/cloudflare/test_cloudflare_rerank_transformation.py From 50b781e8530ec487894e92ee95720f30ad8bfc23 Mon Sep 17 00:00:00 2001 From: prdai Date: Fri, 6 Nov 2026 13:31:13 +0530 Subject: [PATCH 9/9] style(cloudflare): format rebased changes --- litellm/main.py | 1 - .../cloudflare/test_cloudflare_embedding_transformation.py | 3 +-- .../llms/cloudflare/test_cloudflare_rerank_transformation.py | 1 - 3 files changed, 1 insertion(+), 4 deletions(-) diff --git a/litellm/main.py b/litellm/main.py index 05941ea57ba..19d42ec5a7d 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6793,7 +6793,6 @@ def _embedding_headers(param: str, configured: object, model: str, drop_params: ) - def _resolve_vercel_or_cloudflare_embedding_credentials( custom_llm_provider: str, api_base: str | None, diff --git a/tests/unit/llms/cloudflare/test_cloudflare_embedding_transformation.py b/tests/unit/llms/cloudflare/test_cloudflare_embedding_transformation.py index 17ce0f83a18..c17eed9b913 100644 --- a/tests/unit/llms/cloudflare/test_cloudflare_embedding_transformation.py +++ b/tests/unit/llms/cloudflare/test_cloudflare_embedding_transformation.py @@ -4,8 +4,8 @@ from unittest.mock import Mock import pytest import litellm -from litellm.llms.cloudflare.embedding.transformation import CloudflareEmbeddingConfig from litellm.llms.cloudflare.chat.transformation import CloudflareError +from litellm.llms.cloudflare.embedding.transformation import CloudflareEmbeddingConfig from litellm.llms.custom_httpx.http_handler import HTTPHandler from litellm.utils import ProviderConfigManager @@ -21,7 +21,6 @@ class _RecordingHTTPHandler(HTTPHandler): return self.response - def test_provider_config_manager_returns_cloudflare_embedding_config(): config = ProviderConfigManager.get_provider_embedding_config( model="@cf/baai/bge-large-en-v1.5", diff --git a/tests/unit/llms/cloudflare/test_cloudflare_rerank_transformation.py b/tests/unit/llms/cloudflare/test_cloudflare_rerank_transformation.py index ff95603bf2a..d6e63a9df78 100644 --- a/tests/unit/llms/cloudflare/test_cloudflare_rerank_transformation.py +++ b/tests/unit/llms/cloudflare/test_cloudflare_rerank_transformation.py @@ -23,7 +23,6 @@ class _RecordingHTTPHandler(HTTPHandler): return self.response - def test_provider_config_manager_returns_cloudflare_rerank_config(): config = ProviderConfigManager.get_provider_rerank_config( model="@cf/baai/bge-reranker-base",