fix(cloudflare): satisfy basedpyright override and complexity gates

This commit is contained in:
prdai 2026-09-13 15:45:25 +05:30
parent c6a917827d
commit bc21e9dd43
4 changed files with 95 additions and 85 deletions

View file

@ -1,4 +1,5 @@
from collections.abc import Mapping
from typing import Final
import httpx
@ -16,7 +17,7 @@ class CloudflareEmbeddingConfig(VercelAIGatewayEmbeddingConfig):
litellm_params: Mapping[object, object],
stream: bool | None = None,
) -> str:
resolved_base = CloudflareChatConfig.resolve_api_base(api_base).rstrip("/")
resolved_base: Final = CloudflareChatConfig.resolve_api_base(api_base).rstrip("/")
if resolved_base.endswith("/embeddings"):
return resolved_base
return f"{resolved_base}/embeddings"
@ -25,6 +26,6 @@ class CloudflareEmbeddingConfig(VercelAIGatewayEmbeddingConfig):
self,
error_message: str,
status_code: int,
headers: Mapping[object, object] | httpx.Headers,
headers: Mapping[str, str] | httpx.Headers,
) -> CloudflareError:
return CloudflareError(status_code=status_code, message=error_message)

View file

@ -1,6 +1,6 @@
import json
from collections.abc import Mapping, Sequence
from typing import Final
from typing import Final, cast # noqa: TID251 # json boundaries only expose untyped mappings
import httpx
from typing_extensions import NotRequired, ReadOnly, TypedDict
@ -11,7 +11,7 @@ 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 litellm.types.rerank import RerankResponse, RerankResponseDocument, RerankResponseResult
from ..chat.transformation import CloudflareError
@ -22,7 +22,7 @@ class CloudflareRerankContext(TypedDict):
class CloudflareRerankRequest(TypedDict):
query: ReadOnly[str]
contexts: ReadOnly[Sequence[Mapping[str, str]]]
contexts: ReadOnly[Sequence[CloudflareRerankContext]]
top_k: NotRequired[ReadOnly[object]]
@ -41,13 +41,13 @@ class EmptyRequestData(TypedDict, total=False):
pass
EMPTY_REQUEST_DATA: Mapping[str, object] = EmptyRequestData()
EMPTY_REQUEST_DATA: Final[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"))
account_id: Final = 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"
@ -55,14 +55,15 @@ class CloudflareRerankConfig(BaseRerankConfig):
@staticmethod
def _document_to_context(
document: object,
) -> Mapping[str, str]:
) -> CloudflareRerankContext:
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))
document_mapping: Final = cast(Mapping[str, object], document) # cast-ok: isinstance proves only a bare Mapping
text: Final = document_mapping.get("text")
return CloudflareRerankContext(text=text if isinstance(text, str) else json.dumps(document_mapping))
@staticmethod
def _response_items(
@ -75,14 +76,19 @@ class CloudflareRerankConfig(BaseRerankConfig):
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
result: Final = response_json.get("result", response_json)
result_mapping: Final = (
cast(Mapping[str, object], result) # cast-ok: isinstance proves only a bare Mapping
if isinstance(result, Mapping)
else None
)
response: Final = result_mapping.get("response") if result_mapping is not None 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)
return tuple(cast(Sequence[object], response)) # cast-ok: JSON list narrows only to list[Unknown]
@staticmethod
def _transform_response_item(
@ -93,21 +99,22 @@ class CloudflareRerankConfig(BaseRerankConfig):
if not isinstance(item, Mapping):
raise TypeError("Invalid item in Cloudflare rerank response")
index = item.get("id")
score = item.get("score")
item_mapping: Final = cast(Mapping[str, object], item) # cast-ok: isinstance proves only a bare Mapping
index: Final = item_mapping.get("id")
score: Final = item_mapping.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])
document: Final = CloudflareRerankConfig._document_to_context(documents[index])
return RerankResponseResult(
index=index,
relevance_score=float(score),
document=document,
document=RerankResponseDocument(text=document["text"]),
)
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,
@ -115,11 +122,11 @@ class CloudflareRerankConfig(BaseRerankConfig):
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:
resolved_api_key: Final = api_key or get_secret_str("CLOUDFLARE_API_KEY")
if resolved_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}",
"Authorization": f"Bearer {resolved_api_key}",
"accept": "application/json",
"content-type": "application/json",
}
@ -131,8 +138,8 @@ class CloudflareRerankConfig(BaseRerankConfig):
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")
cleaned: Final = (api_base or self._default_api_base()).rstrip("/")
encoded_model: Final = encode_url_path_segments(model, field_name="model")
if cleaned.endswith(f"/{encoded_model}"):
return cleaned
if cleaned.endswith("/ai/v1"):
@ -141,7 +148,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 +159,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,
@ -173,6 +180,7 @@ class CloudflareRerankConfig(BaseRerankConfig):
documents=documents,
)
if top_n is None:
assert return_documents is not None
return CohereRerankParams(
query=query,
documents=documents,
@ -191,26 +199,26 @@ 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],
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")
query: Final = optional_rerank_params.get("query")
raw_documents: Final = 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:
if not isinstance(raw_documents, Sequence) or isinstance(raw_documents, str) or not raw_documents:
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")
documents: Final = cast(Sequence[object], raw_documents) # cast-ok: isinstance proves only a bare Sequence
contexts: Final = tuple(self._document_to_context(document) for document in documents)
top_n: Final = 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,
@ -223,9 +231,9 @@ class CloudflareRerankConfig(BaseRerankConfig):
optional_params: Mapping[str, object] | None = None,
litellm_params: Mapping[str, object] | None = None,
) -> RerankResponse:
request_data = request_data or EMPTY_REQUEST_DATA
resolved_request_data: Final = request_data or EMPTY_REQUEST_DATA
try:
response_json: object = raw_response.json()
response_json: Final = cast(object, raw_response.json()) # cast-ok: httpx Response.json is untyped
except ValueError:
raise CloudflareError(
status_code=raw_response.status_code,
@ -236,33 +244,37 @@ class CloudflareRerankConfig(BaseRerankConfig):
status_code=raw_response.status_code,
message=f"Invalid Cloudflare rerank response: {response_json}",
)
response_data: Final = cast(Mapping[str, object], response_json) # cast-ok: bare Mapping
logging_obj.post_call(
input=request_data.get("query"),
logging_obj.post_call( # pyright: ignore[reportUnknownMemberType] # logging object is legacy-untyped
input=resolved_request_data.get("query"),
api_key=api_key,
additional_args=LoggingAdditionalArgs(complete_input_dict=request_data),
original_response=response_json,
additional_args=LoggingAdditionalArgs(complete_input_dict=resolved_request_data),
original_response=response_data,
)
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(
resolved_optional_params: Final = optional_params or EMPTY_REQUEST_DATA
raw_documents: Final = resolved_optional_params.get("documents")
documents: Final = (
cast(Sequence[object], raw_documents) # cast-ok: isinstance proves only a bare Sequence
if isinstance(raw_documents, Sequence) and not isinstance(raw_documents, str)
else ()
)
return_documents: Final = resolved_optional_params.get("return_documents") is not False
results: Final = tuple(
self._transform_response_item(
item,
documents=documents,
return_documents=return_documents,
)
for item in self._response_items(
response_json=response_json,
response_json=response_data,
status_code=raw_response.status_code,
)
)
return RerankResponse(
id=str(response_json.get("id") or uuid.uuid4()),
results=results,
id=str(response_data.get("id") or uuid.uuid4()),
results=list(results),
)
def get_error_class(

View file

@ -1611,7 +1611,9 @@ class BaseLLMHTTPHandler:
api_key=api_key,
timeout=timeout,
client=client,
optional_rerank_params=optional_rerank_params,
optional_rerank_params=cast( # cast-ok: untyped caller dict
Mapping[str, object], optional_rerank_params
),
)
if client is None or not isinstance(client, HTTPHandler):
@ -1672,6 +1674,7 @@ class BaseLLMHTTPHandler:
except Exception as e:
raise self._handle_error(e=e, provider_config=provider_config)
resolved_optional_rerank_params: Final = dict(optional_rerank_params or {})
return provider_config.transform_rerank_response(
model=model,
raw_response=response,
@ -1679,7 +1682,7 @@ class BaseLLMHTTPHandler:
logging_obj=logging_obj,
api_key=api_key,
request_data=request_data,
optional_params=optional_rerank_params,
optional_params=resolved_optional_rerank_params,
)
def _prepare_decisions_request(

View file

@ -6793,6 +6793,31 @@ def _embedding_headers(param: str, configured: object, model: str, drop_params:
)
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)
@ -7254,41 +7279,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,