litellm/tests/e2e/provider_cache_routing.py
Yuneng Jiang 1972a30def
revert(e2e): unmount Gemini, its api_base means two things
Build 227 mounted Gemini and turned TestGeminiFiles::test_gemini_file_upload
red. litellm's two Gemini endpoints disagree about what api_base means.
Chat composes {api_base}/models/{model}:{endpoint} and defaults api_base to
https://generativelanguage.googleapis.com/v1beta, so the version lives
inside it. File upload composes {api_base}/upload/v1beta/files and defaults
to the host root, so the version lives outside it. A single api_base cannot
satisfy both, and a registration carries no signal about which endpoint the
deployment will be used for, so the edge cannot route one and not the other.

Backing it out rather than working around it. The cache must never turn a
passing test red, which is the same rule the Bedrock model allowlist
follows, and Gemini was 7 of roughly 1030 edge calls in that build. Anyone
pointing litellm's Gemini provider at an AI gateway or a corporate proxy
hits this too, so the fix belongs in litellm; mounting Gemini is one line
once it lands.

This reverts commit 8a553ceb58.
2026-09-16 08:37:25 -07:00

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})