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

feat(cloudflare): add embeddings support

Browse source 
Co-authored-by: Codex <codex@openai.com>
This commit is contained in:
prdai 2026-07-29 10:32:16 +05:30
parent c2c2a623c0
commit 1af08ff00c
8 changed files with 205 additions and 3 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) | ✅ | ✅ | ✅ | ✅ | | | | | | ✅ |

3
litellm/__init__.py
View file

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

5
litellm/_lazy_imports_registry.py
View file

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

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:

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

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

20
litellm/main.py
View file

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

2
litellm/utils.py
View file

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

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

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