mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-21 00:21:49 +00:00
7c2234be3a
The allowlist rejects an unlisted model before the region resolver runs, so the two negative cases that used to cover the resolver were passing for the wrong reason and two mutations of it survived. Answering an env-referenced region with the default mount is only sound because every allowlisted model is a `us.` profile that fans out across the US regions, so assert that on the list itself and drop the per-call branch it made unreachable.
89 lines
3.6 KiB
Python
89 lines
3.6 KiB
Python
from __future__ import annotations
|
|
|
|
from collections.abc import Callable
|
|
from contextvars import ContextVar
|
|
from typing import Final
|
|
|
|
from models import LiteLLMParamsBody, ModelMode
|
|
|
|
LIVE_PROVIDER_REQUIRED: Final[ContextVar[bool]] = ContextVar("live_provider_required", default=False)
|
|
|
|
DEFAULT_BEDROCK_REGION: Final = "us-east-1"
|
|
BEDROCK_CROSS_REGION_PREFIX: Final = "us."
|
|
BEDROCK_EDGE_MODELS: Final = frozenset(
|
|
{
|
|
"us.anthropic.claude-haiku-4-5-20251001-v1:0",
|
|
"us.anthropic.claude-sonnet-4-5-20250929-v1:0",
|
|
"us.anthropic.claude-sonnet-5",
|
|
"us.anthropic.claude-opus-4-7",
|
|
}
|
|
)
|
|
ENV_REFERENCE_PREFIX: Final = "os.environ/"
|
|
|
|
|
|
def bedrock_region(declared: str | None) -> str:
|
|
"""The region whose edge mount a deployment belongs to.
|
|
|
|
Most Bedrock deployments declare `os.environ/AWS_REGION`, which only the
|
|
proxy can resolve from its own environment; the run pod does not share it.
|
|
Answering those with the default mount is correct because every model on the
|
|
edge allowlist is a `us.` inference profile, which fans out across the US
|
|
regions and is reachable from any of them. That invariant is enforced on the
|
|
allowlist itself rather than re-checked per call."""
|
|
if declared is None or declared.startswith(ENV_REFERENCE_PREFIX):
|
|
return DEFAULT_BEDROCK_REGION
|
|
return declared
|
|
|
|
|
|
def bedrock_mount(params: LiteLLMParamsBody) -> str | None:
|
|
"""The edge mount a Bedrock deployment belongs to, or None.
|
|
|
|
The allowlist mirrors the runner role's IAM policy, which names its models
|
|
one by one. A model outside it would be re-signed with an identity that
|
|
cannot invoke it and come back 403 from Bedrock, so an unlisted model keeps
|
|
its direct path and loses only caching. Adding a model is a policy edit in
|
|
litellm-ops and a line here."""
|
|
route: Final = params.model.partition("/")[2]
|
|
model: Final = route.partition("/")[2] or route
|
|
if model not in BEDROCK_EDGE_MODELS:
|
|
return None
|
|
return f"bedrock/{bedrock_region(params.aws_region_name)}"
|
|
|
|
|
|
def route_bedrock(
|
|
params: LiteLLMParamsBody, base_for: Callable[[str], str | None], mode: ModelMode | None,
|
|
) -> LiteLLMParamsBody:
|
|
"""Deployments that carry their own AWS identity stay off the edge. The edge
|
|
re-signs with the run pod's role, so routing an `aws_role_name` deployment
|
|
would quietly replace the very assume-role chain that test exists to prove."""
|
|
if mode is not None or params.aws_role_name is not None or params.aws_access_key_id is not None:
|
|
return params
|
|
if params.api_base is not None or params.aws_bedrock_runtime_endpoint is not None:
|
|
return params
|
|
mount: Final = bedrock_mount(params)
|
|
if mount is None:
|
|
return params
|
|
base: Final = base_for(mount)
|
|
if base is None:
|
|
return params
|
|
return params.model_copy(update={"aws_bedrock_runtime_endpoint": base})
|
|
|
|
|
|
def route_cache_model(
|
|
params: LiteLLMParamsBody, base_for: Callable[[str], str | None], *, enabled: bool, mode: ModelMode | None = None,
|
|
) -> LiteLLMParamsBody:
|
|
if not enabled or LIVE_PROVIDER_REQUIRED.get() or params.mock_response is not None:
|
|
return params
|
|
if params.litellm_credential_name is not None:
|
|
return params
|
|
provider: Final = params.model.partition("/")[0]
|
|
if provider == "bedrock":
|
|
return route_bedrock(params, base_for, mode)
|
|
if mode == "realtime" or params.api_base is not None:
|
|
return params
|
|
if provider not in {"openai", "anthropic"}:
|
|
return params
|
|
base: Final = base_for(provider)
|
|
if base is None:
|
|
return params
|
|
return params.model_copy(update={"api_base": f"{base}/v1" if provider == "openai" else base})
|