mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
fix(cloudflare): satisfy basedpyright override and complexity gates
This commit is contained in:
parent
c6a917827d
commit
bc21e9dd43
4 changed files with 95 additions and 85 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue