mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
fix(cloudflare): satisfy basedpyright override and complexity gates
This commit is contained in:
parent
70c919aece
commit
86acd9d753
2 changed files with 33 additions and 41 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue