fix(cloudflare): satisfy basedpyright override and complexity gates

This commit is contained in:
prdai 2026-09-13 15:45:25 +05:30
parent 70c919aece
commit 86acd9d753
2 changed files with 33 additions and 41 deletions

View file

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

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