From bc21e9dd4324a8b10ebdfc2b0eb9ddea2a945aa2 Mon Sep 17 00:00:00 2001 From: prdai Date: Sun, 13 Sep 2026 15:45:25 +0530 Subject: [PATCH] fix(cloudflare): satisfy basedpyright override and complexity gates --- .../cloudflare/embedding/transformation.py | 5 +- .../llms/cloudflare/rerank/transformation.py | 106 ++++++++++-------- litellm/llms/custom_httpx/llm_http_handler.py | 7 +- litellm/main.py | 62 +++++----- 4 files changed, 95 insertions(+), 85 deletions(-) diff --git a/litellm/llms/cloudflare/embedding/transformation.py b/litellm/llms/cloudflare/embedding/transformation.py index aa62526322e..2d75c8dab98 100644 --- a/litellm/llms/cloudflare/embedding/transformation.py +++ b/litellm/llms/cloudflare/embedding/transformation.py @@ -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) diff --git a/litellm/llms/cloudflare/rerank/transformation.py b/litellm/llms/cloudflare/rerank/transformation.py index 17a5a840bf3..4cc9b4f374a 100644 --- a/litellm/llms/cloudflare/rerank/transformation.py +++ b/litellm/llms/cloudflare/rerank/transformation.py @@ -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( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index 8e7cf857be3..880d5322a1e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -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( diff --git a/litellm/main.py b/litellm/main.py index d6d85ffff60..05941ea57ba 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -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,