mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
8a553ceb58
Gemini needs none of the machinery Bedrock needed. litellm composes
{api_base}/models/{model}:{endpoint} from a custom api_base, so a plain
path-prefixed mount reaches it, and the credential travels as a static
x-goog-api-key header that no host rewrite invalidates. Nothing is
re-signed and nothing leaves the cache key, so a recording still cannot
cross credentials.
A finished turn names a finishReason on every candidate and reports
usageMetadata. The reason is read as a string rather than compared to
STOP: MAX_TOKENS and the safety reasons end a turn just as finally, and
rejecting them would send every one of them upstream forever. Streaming
is the half worth care. Gemini repeats usageMetadata on every chunk and
names a finishReason only on the last, so the terminator is the final
event rather than any event, and a stream the connection cut short ends
on a chunk carrying usage and no reason.
The mount's upstream base carries the API version, so the path the rules
see is /v1beta/models/..., not the one the proxy sent. The first version
of this anchored the rule at the start of that path, which passed every
test against a stub with no version prefix and would have cached nothing
at all in a real run. Caught by replaying the rules over responses
captured from live gemini-2.5-flash, which is also why the tests now
mount their stub under the version prefix.
Vertex stays unmounted and is a separate provider here: litellm grafts
the default Vertex path onto an api_base only when that api_base has no
path of its own, so Vertex needs a root-mounted edge on its own port.
90 lines
3.7 KiB
Python
90 lines
3.7 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/"
|
|
EDGE_PROVIDERS: Final = frozenset({"openai", "anthropic", "gemini"})
|
|
|
|
|
|
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 EDGE_PROVIDERS:
|
|
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})
|