From 1af08ff00c1534a2fbfd449c30dbf156aa2d9cfb Mon Sep 17 00:00:00 2001 From: prdai Date: Wed, 29 Jul 2026 10:32:16 +0530 Subject: [PATCH] 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 901cc5b0cea..4bf1a3b45a8 100644 --- a/README.md +++ b/README.md @@ -292,7 +292,7 @@ For MCP OAuth, an upstream may advertise dynamic client registration but refuse | [Bytez (`bytez`)](https://docs.litellm.ai/docs/providers/bytez) | ✅ | ✅ | ✅ | | | | | | | | | [Cerebras (`cerebras`)](https://docs.litellm.ai/docs/providers/cerebras) | ✅ | ✅ | ✅ | | | | | | | | | [Clarifai (`clarifai`)](https://docs.litellm.ai/docs/providers/clarifai) | ✅ | ✅ | ✅ | | | | | | | | -| [Cloudflare AI Workers (`cloudflare`)](https://docs.litellm.ai/docs/providers/cloudflare_workers) | ✅ | ✅ | ✅ | | | | | | | | +| [Cloudflare AI Workers (`cloudflare`)](https://docs.litellm.ai/docs/providers/cloudflare_workers) | ✅ | ✅ | ✅ | ✅ | | | | | | | | [Codestral (`codestral`)](https://docs.litellm.ai/docs/providers/codestral) | ✅ | ✅ | ✅ | | | | | | | | | [Cognition (`cognition`)](https://docs.litellm.ai/docs/providers/cognition) | ✅ | ✅ | ✅ | | | | | | | | | [Cohere (`cohere`)](https://docs.litellm.ai/docs/providers/cohere) | ✅ | ✅ | ✅ | ✅ | | | | | | ✅ | diff --git a/litellm/__init__.py b/litellm/__init__.py index ccfbf80369f..04e251e0983 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -1643,6 +1643,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 dc323c8cc15..d6dcaa0783a 100644 --- a/litellm/_lazy_imports_registry.py +++ b/litellm/_lazy_imports_registry.py @@ -164,6 +164,7 @@ LLM_CONFIG_NAMES: Final = ( "LlamaAPIConfig", "TogetherAITextCompletionConfig", "CloudflareChatConfig", + "CloudflareEmbeddingConfig", "NovitaConfig", "PetalsConfig", "OllamaChatConfig", @@ -714,6 +715,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 ef11e51a5eb..eedf8169947 100644 --- a/litellm/llms/cloudflare/chat/transformation.py +++ b/litellm/llms/cloudflare/chat/transformation.py @@ -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: 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 17edafcdfca..5a88b25779c 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -6546,6 +6546,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 7732cd88cb5..b6bb76746d1 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8478,6 +8478,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, + )