diff --git a/litellm/llms/cloudflare/rerank/transformation.py b/litellm/llms/cloudflare/rerank/transformation.py index 17a5a840bf3..e52924e1b6b 100644 --- a/litellm/llms/cloudflare/rerank/transformation.py +++ b/litellm/llms/cloudflare/rerank/transformation.py @@ -107,7 +107,7 @@ class CloudflareRerankConfig(BaseRerankConfig): ) 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, @@ -141,7 +141,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 +152,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, @@ -191,7 +191,7 @@ 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], @@ -206,11 +206,10 @@ class CloudflareRerankConfig(BaseRerankConfig): 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) + return CloudflareRerankRequest(query=query, contexts=contexts) + return CloudflareRerankRequest(query=query, contexts=contexts, top_k=top_n) def transform_rerank_response( self, diff --git a/litellm/main.py b/litellm/main.py index 5a88b25779c..d736c8442cc 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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,41 +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, - 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,