scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 74

Merge 86acd9d753 into 252c71c0b2

Browse source
This commit is contained in:
Ranuga 2026-09-15 20:28:26 -04:00 • committed by GitHub
commit e377877051
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 984 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

2
README.md
View file

@ -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) | ✅ | ✅ | ✅ | ✅ | | | | | | ✅ |

6
litellm/__init__.py
View file

@ -1643,6 +1643,12 @@ if TYPE_CHECKING:
from .llms.cloudflare.chat.transformation import (
CloudflareChatConfig as CloudflareChatConfig,
)
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

10
litellm/_lazy_imports_registry.py
View file

@ -164,6 +164,8 @@ LLM_CONFIG_NAMES: Final = (
"LlamaAPIConfig",
"TogetherAITextCompletionConfig",
"CloudflareChatConfig",
"CloudflareEmbeddingConfig",
"CloudflareRerankConfig",
"NovitaConfig",
"PetalsConfig",
"OllamaChatConfig",
@ -714,6 +716,14 @@ _LLM_CONFIGS_IMPORT_MAP: Final = {
".llms.cloudflare.chat.transformation",
"CloudflareChatConfig",
),
"CloudflareEmbeddingConfig": (
".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"),

4
litellm/llms/cloudflare/chat/transformation.py
View file

@ -37,7 +37,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,
@ -46,7 +46,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:

30
litellm/llms/cloudflare/embedding/transformation.py Normal file
View file

@ -0,0 +1,30 @@
from collections.abc import Mapping
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: Mapping[object, object] | httpx.Headers,
) -> CloudflareError:
return CloudflareError(status_code=status_code, message=error_message)

273
litellm/llms/cloudflare/rerank/transformation.py Normal file
View file

@ -0,0 +1,273 @@
import json
from collections.abc import Mapping, Sequence
from typing import Final
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[str | Mapping[str, object]]]
top_n: NotRequired[ReadOnly[int]]
return_documents: NotRequired[ReadOnly[bool]]
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( # pyright: ignore[reportIncompatibleMethodOverride] # base annotates the return as dict; the read-only Mapping is intentional
self,
headers: Mapping[str, object],
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:
raise ValueError("Missing Cloudflare API Key - set CLOUDFLARE_API_KEY or pass api_key explicitly")
default_headers: Final = {
"Authorization": f"Bearer {api_key}",
"accept": "application/json",
"content-type": "application/json",
}
return {**default_headers, **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( # pyright: ignore[reportIncompatibleMethodOverride] # base annotates the return as list; a read-only tuple is intentional
self,
model: str,
) -> Sequence[str]:
return (
"query",
"documents",
"top_n",
"return_documents",
)
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,
drop_params: bool,
query: str,
documents: Sequence[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( # 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")
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)
top_n = optional_rerank_params.get("top_n")
if top_n is None:
return CloudflareRerankRequest(query=query, contexts=contexts)
return CloudflareRerankRequest(query=query, contexts=contexts, 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: Mapping[str, object] | httpx.Headers,
) -> BaseLLMException:
return CloudflareError(status_code=status_code, message=error_message)

4
litellm/llms/custom_httpx/llm_http_handler.py
View file

@ -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(

41
litellm/main.py
View file

@ -6071,6 +6071,30 @@ async def aembedding(*args, **kwargs) -> EmbeddingResponse:
)
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)
@ -6531,21 +6555,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,

3
litellm/rerank_api/main.py
View file

@ -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,

6
litellm/utils.py
View file

@ -8469,6 +8469,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:
@ -8484,7 +8486,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,
@ -8531,6 +8533,8 @@ class ProviderConfigManager:
)
return get_dashscope_family_rerank_config(provider.value)
elif litellm.LlmProviders.CLOUDFLARE == provider:
return litellm.CloudflareRerankConfig()
return litellm.CohereRerankConfig()
@staticmethod

4
provider_endpoints_support.json
View file

@ -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
}

56
tests/local_testing/test_embedding.py
View file

@ -1026,6 +1026,62 @@ def test_hosted_vllm_embedding(monkeypatch):
assert json_data["model"] == "jina-embeddings-v3"
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_cloudflare_embedding_dispatch(monkeypatch):
monkeypatch.setattr(litellm, "cloudflare_api_key", None)
monkeypatch.setattr(litellm, "api_key", None)
monkeypatch.setattr(litellm, "api_base", None)
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 = MagicMock()
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)
client = _RecordingHTTPHandler(raw_response)
embedding(
model="cloudflare/@cf/baai/bge-large-en-v1.5",
input=["Hello world"],
api_key="cf-key",
api_base="https://example.com/ai/v1",
client=client,
caching=False,
)
request = client.requests[0]
assert request["url"] == "https://example.com/ai/v1/embeddings"
assert request["headers"]["Authorization"] == "Bearer cf-key"
def test_cloudflare_embedding_dispatch_requires_api_key(monkeypatch):
monkeypatch.setattr(litellm, "cloudflare_api_key", None)
monkeypatch.setattr(litellm, "api_key", None)
monkeypatch.delenv("CLOUDFLARE_API_KEY", raising=False)
with pytest.raises(litellm.APIConnectionError, match="Missing Cloudflare API Key"):
embedding(
model="cloudflare/@cf/baai/bge-large-en-v1.5",
input=["Hello world"],
caching=False,
)
def test_llamafile_embedding(monkeypatch):
monkeypatch.setenv("LLAMAFILE_API_BASE", "http://localhost:8080/v1")
from litellm.llms.custom_httpx.http_handler import HTTPHandler

152
tests/test_litellm/llms/cloudflare/test_cloudflare_embedding_transformation.py Normal file
View file

@ -0,0 +1,152 @@
import json
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.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",
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")
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)
client = _RecordingHTTPHandler(raw_response)
response = litellm.embedding(
model="cloudflare/@cf/baai/bge-large-en-v1.5",
input=["hello"],
api_key="cf-key",
client=client,
caching=False,
)
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"
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,
)

414
tests/test_litellm/llms/cloudflare/test_cloudflare_rerank_transformation.py Normal file
View file

@ -0,0 +1,414 @@
import json
from unittest.mock import MagicMock, Mock
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
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",
provider=litellm.LlmProviders.CLOUDFLARE,
api_base=None,
present_version_params=[],
)
assert isinstance(config, CloudflareRerankConfig)
def test_get_complete_url_uses_native_workers_ai_endpoint(monkeypatch):
monkeypatch.setenv("CLOUDFLARE_ACCOUNT_ID", "account-id")
config = CloudflareRerankConfig()
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(monkeypatch):
monkeypatch.delenv("CLOUDFLARE_ACCOUNT_ID", raising=False)
config = CloudflareRerankConfig()
with pytest.raises(ValueError, match="CLOUDFLARE_ACCOUNT_ID"):
config.get_complete_url(None, "@cf/baai/bge-reranker-base")
@pytest.mark.parametrize(
"model,error_match",
(
("../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, error_match):
config = CloudflareRerankConfig()
with pytest.raises(ValueError, match=error_match):
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(monkeypatch):
monkeypatch.delenv("CLOUDFLARE_API_KEY", raising=False)
config = CloudflareRerankConfig()
with 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,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, error_match):
config = CloudflareRerankConfig()
with pytest.raises(ValueError, match=error_match):
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():
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)
client = _RecordingHTTPHandler(raw_response)
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 = 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"]) == {
"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."},
}
]