mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
fix(passthrough): resolve vertex live credentials from db model deployments
The /vertex_ai/live WebSocket passthrough only ever looked at default_vertex_config and the DEFAULT_VERTEXAI_* env vars, so a proxy whose Vertex credentials live in the DB as a model entry with use_in_pass_through had nothing to authenticate with. The upgrade still succeeded and the socket then closed with a bare 1000 on the first client frame, which gave the client no way to tell a misconfiguration from a normal end of session. Credentials now also resolve from the router deployments flagged use_in_pass_through, preferring the one matching the requested model, and a failure to mint an access token closes 1011 with a reason naming both ways to configure it. Upstream closes other than a plain 1000 are relayed to the client with their code and reason, so Google's own errors reach the caller. The setup frame's model is rewritten to the full projects/.../publishers/google/models resource path, which is what Vertex expects and what lets a bare model id or a gateway alias work over this route.
This commit is contained in:
mateo-berri
parent
5290150a05
commit
021a09b156
7 changed files with 701 additions and 42 deletions
|
|
@ -244,6 +244,9 @@ AIOHTTP_NEEDS_CLEANUP_CLOSED: Final = (3, 13, 0) <= sys.version_info < (
|
|||
_max_size_env: Final = os.getenv("REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES")
|
||||
REALTIME_WEBSOCKET_MAX_MESSAGE_SIZE_BYTES: Final = int(_max_size_env) if _max_size_env is not None else None
|
||||
|
||||
# RFC 6455 caps the close frame payload at 125 bytes, 2 of which carry the status code
|
||||
WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123
|
||||
|
||||
# SSL/TLS cipher configuration for faster handshakes
|
||||
# Strategy: Strongly prefer fast modern ciphers, but allow fallback to commonly supported ones
|
||||
# This balances performance with broad compatibility
|
||||
|
|
|
|||
|
|
@ -9,8 +9,9 @@ Use litellm with Anthropic SDK, Vertex AI SDK, Cohere SDK, etc.
|
|||
import json
|
||||
import os
|
||||
import re
|
||||
from collections.abc import Callable
|
||||
from types import MappingProxyType
|
||||
from typing import Annotated, Any, Final, cast
|
||||
from typing import TYPE_CHECKING, Annotated, Any, Final, cast
|
||||
|
||||
import httpx
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket
|
||||
|
|
@ -55,11 +56,15 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
|||
LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
|
||||
LITELLM_PASS_THROUGH_RAW_BODY_STATE_KEY,
|
||||
)
|
||||
from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials
|
||||
from litellm.types.utils import LlmProviders
|
||||
from litellm.utils import ProviderConfigManager
|
||||
|
||||
from .passthrough_endpoint_router import PassthroughEndpointRouter
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router
|
||||
|
||||
vertex_llm_base: Final = VertexBase()
|
||||
router: Final = APIRouter()
|
||||
openai_passthrough_router: Final = APIRouter()
|
||||
|
|
@ -2373,6 +2378,89 @@ async def cursor_proxy_route(
|
|||
return received_value
|
||||
|
||||
|
||||
VERTEX_LIVE_UNCONFIGURED_CLOSE_REASON: Final = (
|
||||
"Vertex AI auth failed: set a use_in_pass_through vertex model, default_vertex_config, or DEFAULT_VERTEXAI_* env"
|
||||
)
|
||||
|
||||
VERTEX_PUBLISHER_MODEL_PREFIX: Final = "publishers/google/models/"
|
||||
|
||||
|
||||
def _get_llm_router() -> "Router | None":
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
return llm_router
|
||||
|
||||
|
||||
def _resolve_vertex_live_credentials(
|
||||
vertex_project: str | None,
|
||||
vertex_location: str | None,
|
||||
model: str | None,
|
||||
) -> VertexPassThroughCredentials | None:
|
||||
"""
|
||||
Resolution order: credentials registered for the requested project/location, then any DB model entry
|
||||
flagged ``use_in_pass_through``, then ``default_vertex_config`` and the ``DEFAULT_VERTEXAI_*`` env vars
|
||||
"""
|
||||
keyed: Final = passthrough_endpoint_router.get_vertex_credentials(
|
||||
project_id=vertex_project,
|
||||
location=vertex_location,
|
||||
)
|
||||
if keyed is not None and keyed.vertex_project is not None:
|
||||
return keyed
|
||||
from_deployments: Final = passthrough_endpoint_router.get_vertex_credentials_from_router_deployments(model=model)
|
||||
if from_deployments is not None:
|
||||
return from_deployments
|
||||
if keyed is not None:
|
||||
return keyed
|
||||
passthrough_endpoint_router.set_default_vertex_config()
|
||||
return passthrough_endpoint_router.get_vertex_credentials(
|
||||
project_id=vertex_project,
|
||||
location=vertex_location,
|
||||
)
|
||||
|
||||
|
||||
def _build_vertex_live_setup_model_rewriter(
|
||||
vertex_project: str | None,
|
||||
vertex_location: str | None,
|
||||
llm_router: "Router | None",
|
||||
) -> Callable[[str], str] | None:
|
||||
"""
|
||||
Rewrite the ``setup`` frame's model into the full Vertex resource path the Live API requires.
|
||||
|
||||
Clients address the gateway the way they address LiteLLM (bare id or model alias); Vertex reads anything
|
||||
that is not a ``projects/...`` path as a project name and closes the socket
|
||||
"""
|
||||
if vertex_project is None or vertex_location is None:
|
||||
return None
|
||||
|
||||
def rewrite(setup_model: str) -> str:
|
||||
if setup_model.startswith("projects/"):
|
||||
return setup_model
|
||||
aliased: Final = _resolve_alias_to_upstream_model(setup_model, llm_router)
|
||||
return (
|
||||
f"projects/{vertex_project}/locations/{vertex_location}/"
|
||||
f"{VERTEX_PUBLISHER_MODEL_PREFIX}{aliased.removeprefix(VERTEX_PUBLISHER_MODEL_PREFIX)}"
|
||||
)
|
||||
|
||||
return rewrite
|
||||
|
||||
|
||||
def _resolve_alias_to_upstream_model(setup_model: str, llm_router: "Router | None") -> str:
|
||||
if llm_router is None:
|
||||
return setup_model
|
||||
upstream: Final = next(
|
||||
(
|
||||
deployment["litellm_params"]["model"]
|
||||
for deployment in (llm_router.get_model_list() or ())
|
||||
if deployment.get("model_name") == setup_model
|
||||
),
|
||||
None,
|
||||
)
|
||||
if upstream is None:
|
||||
return setup_model
|
||||
_, provider, _, _ = litellm.get_llm_provider(model=upstream)
|
||||
return upstream.removeprefix(f"{provider}/")
|
||||
|
||||
|
||||
async def vertex_ai_live_websocket_passthrough(
|
||||
websocket: WebSocket,
|
||||
model: str | None = None,
|
||||
|
|
@ -2396,51 +2484,40 @@ async def vertex_ai_live_websocket_passthrough(
|
|||
await websocket.accept()
|
||||
|
||||
incoming_headers: Final = dict(websocket.headers)
|
||||
vertex_credentials_config = passthrough_endpoint_router.get_vertex_credentials(
|
||||
project_id=vertex_project,
|
||||
location=vertex_location,
|
||||
vertex_credentials_config: Final = _resolve_vertex_live_credentials(
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
model=model,
|
||||
)
|
||||
|
||||
if vertex_credentials_config is None:
|
||||
# Attempt to load defaults from environment/config if not already initialised
|
||||
passthrough_endpoint_router.set_default_vertex_config()
|
||||
vertex_credentials_config = passthrough_endpoint_router.get_vertex_credentials(
|
||||
project_id=vertex_project,
|
||||
location=vertex_location,
|
||||
)
|
||||
|
||||
resolved_project = vertex_project
|
||||
resolved_location: str | None = vertex_location
|
||||
credentials_value: str | None = None
|
||||
|
||||
if vertex_credentials_config is not None:
|
||||
resolved_project = resolved_project or vertex_credentials_config.vertex_project
|
||||
temp_location: Final = resolved_location or vertex_credentials_config.vertex_location
|
||||
# Ensure resolved_location is a string
|
||||
if isinstance(temp_location, dict) or temp_location is not None:
|
||||
resolved_location = str(temp_location)
|
||||
else:
|
||||
resolved_location = None
|
||||
credentials_value = (
|
||||
str(vertex_credentials_config.vertex_credentials)
|
||||
if vertex_credentials_config.vertex_credentials is not None
|
||||
else None
|
||||
)
|
||||
configured_project: Final = vertex_project or (
|
||||
vertex_credentials_config.vertex_project if vertex_credentials_config is not None else None
|
||||
)
|
||||
configured_location: Final = vertex_location or (
|
||||
vertex_credentials_config.vertex_location if vertex_credentials_config is not None else None
|
||||
)
|
||||
credentials_value: Final = (
|
||||
str(vertex_credentials_config.vertex_credentials)
|
||||
if vertex_credentials_config is not None and vertex_credentials_config.vertex_credentials is not None
|
||||
else None
|
||||
)
|
||||
|
||||
try:
|
||||
resolved_location = resolved_location or (vertex_llm_base.get_default_vertex_location())
|
||||
if model:
|
||||
resolved_location = vertex_llm_base.get_vertex_region(
|
||||
vertex_region=resolved_location,
|
||||
resolved_location: Final = (
|
||||
vertex_llm_base.get_vertex_region(
|
||||
vertex_region=configured_location or vertex_llm_base.get_default_vertex_location(),
|
||||
model=model,
|
||||
)
|
||||
if model
|
||||
else configured_location or vertex_llm_base.get_default_vertex_location()
|
||||
)
|
||||
|
||||
(
|
||||
access_token,
|
||||
resolved_project,
|
||||
) = await vertex_llm_base._ensure_access_token_async(
|
||||
credentials=credentials_value,
|
||||
project_id=resolved_project,
|
||||
project_id=configured_project,
|
||||
custom_llm_provider="vertex_ai_beta",
|
||||
)
|
||||
except Exception as e:
|
||||
|
|
@ -2453,7 +2530,7 @@ async def vertex_ai_live_websocket_passthrough(
|
|||
request_data={},
|
||||
)
|
||||
if websocket.client_state != WebSocketState.DISCONNECTED:
|
||||
await websocket.close(code=1011, reason="Vertex AI authentication failed")
|
||||
await websocket.close(code=1011, reason=VERTEX_LIVE_UNCONFIGURED_CLOSE_REASON)
|
||||
return
|
||||
|
||||
host_location: Final = resolved_location or vertex_llm_base.get_default_vertex_location()
|
||||
|
|
@ -2485,6 +2562,11 @@ async def vertex_ai_live_websocket_passthrough(
|
|||
forward_headers=False,
|
||||
endpoint="/vertex_ai/live",
|
||||
accept_websocket=False,
|
||||
setup_model_rewriter=_build_vertex_live_setup_model_rewriter(
|
||||
vertex_project=resolved_project,
|
||||
vertex_location=resolved_location,
|
||||
llm_router=_get_llm_router(),
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -5,7 +5,7 @@ import json
|
|||
import posixpath
|
||||
import traceback
|
||||
from base64 import b64encode
|
||||
from collections.abc import AsyncGenerator, Callable, Mapping
|
||||
from collections.abc import AsyncGenerator, Callable, Iterable, Mapping
|
||||
from datetime import datetime
|
||||
from itertools import groupby
|
||||
from typing import Any, Final, TypedDict, cast
|
||||
|
|
@ -32,11 +32,15 @@ from websockets.exceptions import (
|
|||
ConnectionClosedOK,
|
||||
InvalidStatus,
|
||||
)
|
||||
from websockets.frames import Close
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG
|
||||
from litellm.constants import (
|
||||
MAXIMUM_TRACEBACK_LINES_TO_LOG,
|
||||
WEBSOCKET_CLOSE_REASON_MAX_BYTES,
|
||||
)
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.litellm_core_utils.core_helpers import (
|
||||
|
|
@ -1890,6 +1894,52 @@ def create_websocket_passthrough_route(
|
|||
return websocket_endpoint_func
|
||||
|
||||
|
||||
def _rewrite_vertex_live_setup_model(text_data: str, setup_model_rewriter: Callable[[str], str] | None) -> str:
|
||||
"""
|
||||
Rewrite the model of a Vertex AI Live ``setup`` frame, leaving every other frame byte-identical
|
||||
"""
|
||||
if setup_model_rewriter is None:
|
||||
return text_data
|
||||
try:
|
||||
message: Final = json.loads(text_data)
|
||||
except json.JSONDecodeError:
|
||||
return text_data
|
||||
if not isinstance(message, dict):
|
||||
return text_data
|
||||
setup: Final = message.get("setup")
|
||||
if not isinstance(setup, dict):
|
||||
return text_data
|
||||
setup_model: Final = setup.get("model")
|
||||
if not isinstance(setup_model, str):
|
||||
return text_data
|
||||
rewritten_model: Final = setup_model_rewriter(setup_model)
|
||||
if rewritten_model == setup_model:
|
||||
return text_data
|
||||
return json.dumps({**message, "setup": {**setup, "model": rewritten_model}}) # mutable-ok: one-shot json payload
|
||||
|
||||
|
||||
def _truncated_close_reason(reason: str) -> str:
|
||||
"""
|
||||
Fit a close reason inside the byte budget a WebSocket close frame allows, without splitting a character
|
||||
"""
|
||||
encoded: Final = reason.encode("utf-8")
|
||||
if len(encoded) <= WEBSOCKET_CLOSE_REASON_MAX_BYTES:
|
||||
return reason
|
||||
return encoded[:WEBSOCKET_CLOSE_REASON_MAX_BYTES].decode("utf-8", errors="ignore")
|
||||
|
||||
|
||||
def _upstream_close_to_relay(task_results: Iterable[object]) -> Close | None:
|
||||
"""
|
||||
The upstream close worth telling the client about: anything other than a plain, reasonless normal close
|
||||
"""
|
||||
upstream_close: Final = next((result for result in task_results if isinstance(result, Close)), None)
|
||||
if upstream_close is None:
|
||||
return None
|
||||
if upstream_close.code == 1000 and upstream_close.reason == "":
|
||||
return None
|
||||
return upstream_close
|
||||
|
||||
|
||||
async def websocket_passthrough_request(
|
||||
websocket: WebSocket,
|
||||
target: str,
|
||||
|
|
@ -1899,6 +1949,7 @@ async def websocket_passthrough_request(
|
|||
endpoint: str | None = None,
|
||||
cost_per_request: float | None = None,
|
||||
accept_websocket: bool = True,
|
||||
setup_model_rewriter: Callable[[str], str] | None = None,
|
||||
):
|
||||
"""
|
||||
WebSocket passthrough request handler.
|
||||
|
|
@ -1911,6 +1962,7 @@ async def websocket_passthrough_request(
|
|||
forward_headers: Whether to forward incoming headers
|
||||
endpoint: The endpoint path (for logging purposes)
|
||||
cost_per_request: Optional field - cost per request to the target endpoint
|
||||
setup_model_rewriter: Optional rewrite of the setup frame's model before it reaches the upstream
|
||||
"""
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.proxy.proxy_server import proxy_logging_obj
|
||||
|
|
@ -2100,7 +2152,7 @@ async def websocket_passthrough_request(
|
|||
)
|
||||
# Not a JSON message or doesn't contain setup data
|
||||
|
||||
await upstream_ws.send(text_data)
|
||||
await upstream_ws.send(_rewrite_vertex_live_setup_model(text_data, setup_model_rewriter))
|
||||
elif bytes_data is not None:
|
||||
await upstream_ws.send(bytes_data)
|
||||
except asyncio.CancelledError:
|
||||
|
|
@ -2111,8 +2163,8 @@ async def websocket_passthrough_request(
|
|||
)
|
||||
await upstream_ws.close()
|
||||
|
||||
async def forward_upstream_to_client() -> None:
|
||||
"""Forward messages from upstream to client WebSocket"""
|
||||
async def forward_upstream_to_client() -> Close | None:
|
||||
"""Forward messages from upstream to client WebSocket, returning the upstream's close frame"""
|
||||
try:
|
||||
# Wait for the first response from upstream
|
||||
raw_response = await upstream_ws.recv(decode=False)
|
||||
|
|
@ -2177,6 +2229,7 @@ async def websocket_passthrough_request(
|
|||
|
||||
except (ConnectionClosedOK, ConnectionClosedError) as e:
|
||||
verbose_proxy_logger.debug("Upstream WebSocket connection closed: %s", e)
|
||||
return e.rcvd
|
||||
except asyncio.CancelledError:
|
||||
verbose_proxy_logger.debug("asyncio.CancelledError in forward_upstream_to_client")
|
||||
raise
|
||||
|
|
@ -2209,6 +2262,13 @@ async def websocket_passthrough_request(
|
|||
if exception is not None:
|
||||
raise exception
|
||||
|
||||
upstream_close: Final = _upstream_close_to_relay(task.result() for task in done)
|
||||
if upstream_close is not None and websocket.application_state != WebSocketState.DISCONNECTED:
|
||||
await websocket.close(
|
||||
code=upstream_close.code,
|
||||
reason=_truncated_close_reason(upstream_close.reason),
|
||||
)
|
||||
|
||||
end_time: Final = datetime.now()
|
||||
|
||||
# Update passthrough logging payload with response data
|
||||
|
|
@ -2325,7 +2385,10 @@ async def websocket_passthrough_request(
|
|||
if websocket.client_state != WebSocketState.DISCONNECTED:
|
||||
await websocket.close(code=1011, reason="WebSocket passthrough error")
|
||||
finally:
|
||||
if websocket.client_state != WebSocketState.DISCONNECTED:
|
||||
if (
|
||||
websocket.client_state != WebSocketState.DISCONNECTED
|
||||
and websocket.application_state != WebSocketState.DISCONNECTED
|
||||
):
|
||||
await websocket.close()
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
|||
from litellm.secret_managers.main import get_secret_str
|
||||
from litellm.types.llms.vertex_ai import VERTEX_CREDENTIALS_TYPES
|
||||
from litellm.types.passthrough_endpoints.vertex_ai import VertexPassThroughCredentials
|
||||
from litellm.types.router import LiteLLMParamsTypedDict
|
||||
from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.router import Router
|
||||
|
|
@ -120,6 +120,69 @@ class PassthroughEndpointRouter:
|
|||
return None
|
||||
return provider
|
||||
|
||||
def get_vertex_credentials_from_router_deployments(self, model: str | None) -> VertexPassThroughCredentials | None:
|
||||
"""
|
||||
Resolve vertex pass-through credentials from the live router deployments flagged ``use_in_pass_through``.
|
||||
|
||||
``deployment_key_to_vertex_credentials`` is only reachable when the caller names a project and location,
|
||||
which WebSocket clients never do, so DB-stored deployments need this lookup to be usable at all.
|
||||
"""
|
||||
llm_router: Final = self.llm_router_getter()
|
||||
if llm_router is None:
|
||||
return None
|
||||
resolved: Final = tuple(
|
||||
(deployment, credentials)
|
||||
for deployment in (llm_router.get_model_list() or ())
|
||||
if (credentials := self._resolve_vertex_deployment_credentials(deployment["litellm_params"])) is not None
|
||||
)
|
||||
if len(resolved) == 0:
|
||||
return None
|
||||
return next(
|
||||
(
|
||||
credentials
|
||||
for deployment, credentials in resolved
|
||||
if model is not None and self._deployment_matches_model(deployment, model)
|
||||
),
|
||||
resolved[0][1],
|
||||
)
|
||||
|
||||
def _resolve_vertex_deployment_credentials(
|
||||
self, litellm_params: LiteLLMParamsTypedDict
|
||||
) -> VertexPassThroughCredentials | None:
|
||||
if litellm_params.get("use_in_pass_through") is not True:
|
||||
return None
|
||||
if self._get_deployment_provider(litellm_params) != "vertex_ai":
|
||||
return None
|
||||
credential_name: Final = litellm_params.get("litellm_credential_name")
|
||||
credential_values: Final = (
|
||||
CredentialAccessor.get_credential_values(credential_name) if credential_name is not None else None
|
||||
)
|
||||
vertex_project: Final = _get_str_value(credential_values, "vertex_project") or litellm_params.get(
|
||||
"vertex_project"
|
||||
)
|
||||
vertex_location: Final = _get_str_value(credential_values, "vertex_location") or litellm_params.get(
|
||||
"vertex_location"
|
||||
)
|
||||
vertex_credentials: Final = _get_str_value(credential_values, "vertex_credentials") or litellm_params.get(
|
||||
"vertex_credentials"
|
||||
)
|
||||
if vertex_project is None or vertex_location is None:
|
||||
return None
|
||||
return VertexPassThroughCredentials(
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
vertex_credentials=vertex_credentials,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _deployment_matches_model(deployment: DeploymentTypedDict, model: str) -> bool:
|
||||
upstream_model: Final = deployment["litellm_params"].get("model")
|
||||
return model in (
|
||||
deployment.get("model_name"),
|
||||
upstream_model,
|
||||
upstream_model.split("/", 1)[-1] if upstream_model is not None else None,
|
||||
)
|
||||
|
||||
def _get_vertex_env_vars(self) -> VertexPassThroughCredentials:
|
||||
"""
|
||||
Helper to get vertex pass through config from environment variables
|
||||
|
|
|
|||
|
|
@ -3887,3 +3887,104 @@ class TestComprehendMedicalProxyRoute:
|
|||
user_api_key_dict=Mock(),
|
||||
)
|
||||
assert exc_info.value.status_code == 400
|
||||
|
||||
|
||||
class TestVertexAILiveWebsocketPassthrough:
|
||||
def _websocket(self):
|
||||
from starlette.websockets import WebSocketState
|
||||
|
||||
websocket = MagicMock()
|
||||
websocket.accept = AsyncMock()
|
||||
websocket.close = AsyncMock()
|
||||
websocket.headers = {}
|
||||
websocket.client_state = WebSocketState.CONNECTED
|
||||
return websocket
|
||||
|
||||
def _clear_vertex_env(self, monkeypatch):
|
||||
monkeypatch.delenv("DEFAULT_VERTEXAI_PROJECT", raising=False)
|
||||
monkeypatch.delenv("DEFAULT_VERTEXAI_LOCATION", raising=False)
|
||||
monkeypatch.delenv("DEFAULT_GOOGLE_APPLICATION_CREDENTIALS", raising=False)
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_uses_db_deployment_credentials_without_query_params(self, monkeypatch):
|
||||
from litellm.proxy.pass_through_endpoints import (
|
||||
llm_passthrough_endpoints as passthrough_module,
|
||||
)
|
||||
|
||||
llm_router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gemini-live",
|
||||
"litellm_params": {
|
||||
"model": "vertex_ai/gemini-live-2.5-flash",
|
||||
"use_in_pass_through": True,
|
||||
"vertex_project": "proj-db",
|
||||
"vertex_location": "global",
|
||||
"vertex_credentials": '{"type": "service_account"}',
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router)
|
||||
monkeypatch.setattr(
|
||||
passthrough_module.passthrough_endpoint_router, "default_vertex_config", None
|
||||
)
|
||||
self._clear_vertex_env(monkeypatch)
|
||||
websocket = self._websocket()
|
||||
ensure_token = AsyncMock(return_value=("token-abc", "proj-db"))
|
||||
ws_passthrough = AsyncMock()
|
||||
|
||||
with (
|
||||
patch.object(passthrough_module.vertex_llm_base, "_ensure_access_token_async", ensure_token),
|
||||
patch.object(passthrough_module, "websocket_passthrough_request", ws_passthrough),
|
||||
):
|
||||
await passthrough_module.vertex_ai_live_websocket_passthrough(
|
||||
websocket=websocket,
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
|
||||
ensure_token.assert_awaited_once_with(
|
||||
credentials='{"type": "service_account"}',
|
||||
project_id="proj-db",
|
||||
custom_llm_provider="vertex_ai_beta",
|
||||
)
|
||||
passthrough_kwargs = ws_passthrough.await_args.kwargs
|
||||
assert passthrough_kwargs["target"] == (
|
||||
"wss://aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent"
|
||||
)
|
||||
assert passthrough_kwargs["custom_headers"]["Authorization"] == "Bearer token-abc"
|
||||
rewriter = passthrough_kwargs["setup_model_rewriter"]
|
||||
assert rewriter("gemini-live") == (
|
||||
"projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash"
|
||||
)
|
||||
websocket.close.assert_not_awaited()
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_credential_failure_close_names_configuration_options(self, monkeypatch):
|
||||
from litellm.proxy.pass_through_endpoints import (
|
||||
llm_passthrough_endpoints as passthrough_module,
|
||||
)
|
||||
|
||||
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
|
||||
monkeypatch.setattr(
|
||||
passthrough_module.passthrough_endpoint_router, "default_vertex_config", None
|
||||
)
|
||||
self._clear_vertex_env(monkeypatch)
|
||||
websocket = self._websocket()
|
||||
ensure_token = AsyncMock(side_effect=Exception("Unable to find your credentials"))
|
||||
|
||||
with (
|
||||
patch.object(passthrough_module.vertex_llm_base, "_ensure_access_token_async", ensure_token),
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging,
|
||||
):
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
await passthrough_module.vertex_ai_live_websocket_passthrough(
|
||||
websocket=websocket,
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
|
||||
close_kwargs = websocket.close.await_args.kwargs
|
||||
assert close_kwargs["code"] == 1011
|
||||
assert "use_in_pass_through" in close_kwargs["reason"]
|
||||
assert "default_vertex_config" in close_kwargs["reason"]
|
||||
assert len(close_kwargs["reason"].encode("utf-8")) <= 123
|
||||
|
|
|
|||
|
|
@ -4994,6 +4994,212 @@ async def test_websocket_passthrough_forwards_non_ascii_first_frame():
|
|||
assert all(call.kwargs.get("code") != 1011 for call in websocket.close.await_args_list)
|
||||
|
||||
|
||||
class ClosingUpstreamWebSocket:
|
||||
def __init__(self, close_exc: Exception):
|
||||
self._close_exc = close_exc
|
||||
self.close = AsyncMock()
|
||||
self.send = AsyncMock()
|
||||
|
||||
async def recv(self, decode: bool = True):
|
||||
raise self._close_exc
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
raise StopAsyncIteration
|
||||
|
||||
|
||||
class RecordingUpstreamWebSocket:
|
||||
def __init__(self):
|
||||
self.close = AsyncMock()
|
||||
self.send = AsyncMock()
|
||||
|
||||
async def recv(self, decode: bool = True):
|
||||
await asyncio.Event().wait()
|
||||
|
||||
def __aiter__(self):
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
raise StopAsyncIteration
|
||||
|
||||
|
||||
def _client_websocket(receive):
|
||||
from starlette.websockets import WebSocketState
|
||||
|
||||
websocket = MagicMock()
|
||||
websocket.accept = AsyncMock()
|
||||
websocket.send_text = AsyncMock()
|
||||
websocket.send_bytes = AsyncMock()
|
||||
websocket.receive = receive
|
||||
websocket.headers = {}
|
||||
websocket.client_state = WebSocketState.CONNECTED
|
||||
websocket.application_state = WebSocketState.CONNECTED
|
||||
|
||||
def _mark_closed(*args, **kwargs):
|
||||
websocket.application_state = WebSocketState.DISCONNECTED
|
||||
|
||||
websocket.close = AsyncMock(side_effect=_mark_closed)
|
||||
return websocket
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _patched_websocket_passthrough_environment(upstream_ws):
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.proxy_logging_obj") as mock_proxy_logging,
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.connect",
|
||||
return_value=FakeUpstreamConnect(upstream_ws),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.pass_through_endpoints.GLOBAL_LOGGING_WORKER"
|
||||
) as mock_worker,
|
||||
):
|
||||
mock_proxy_logging.pre_call_hook = AsyncMock(return_value={})
|
||||
mock_proxy_logging.post_call_success_hook = AsyncMock()
|
||||
mock_proxy_logging.post_call_failure_hook = AsyncMock()
|
||||
mock_worker.ensure_initialized_and_enqueue = MagicMock(
|
||||
side_effect=lambda async_coroutine: async_coroutine.close()
|
||||
)
|
||||
yield
|
||||
|
||||
|
||||
async def _pending_receive():
|
||||
await asyncio.Event().wait()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_relays_upstream_policy_close_to_client():
|
||||
from websockets.exceptions import ConnectionClosedError
|
||||
from websockets.frames import Close
|
||||
|
||||
upstream_reason = "Publisher Model `projects/p/locations/global/publishers/google/models/nope` was not found"
|
||||
upstream_ws = ClosingUpstreamWebSocket(
|
||||
ConnectionClosedError(
|
||||
rcvd=Close(1008, upstream_reason),
|
||||
sent=Close(1008, ""),
|
||||
rcvd_then_sent=True,
|
||||
)
|
||||
)
|
||||
websocket = _client_websocket(_pending_receive)
|
||||
|
||||
with _patched_websocket_passthrough_environment(upstream_ws):
|
||||
await websocket_passthrough_request(
|
||||
websocket=websocket,
|
||||
target="wss://aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent",
|
||||
custom_headers={"Authorization": "Bearer token"},
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
forward_headers=False,
|
||||
endpoint="/vertex_ai/live",
|
||||
accept_websocket=False,
|
||||
)
|
||||
|
||||
websocket.close.assert_awaited_once_with(code=1008, reason=upstream_reason)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_keeps_normal_upstream_close_normal():
|
||||
from websockets.exceptions import ConnectionClosedOK
|
||||
from websockets.frames import Close
|
||||
|
||||
upstream_ws = ClosingUpstreamWebSocket(
|
||||
ConnectionClosedOK(rcvd=Close(1000, ""), sent=Close(1000, ""), rcvd_then_sent=True)
|
||||
)
|
||||
websocket = _client_websocket(_pending_receive)
|
||||
|
||||
with _patched_websocket_passthrough_environment(upstream_ws):
|
||||
await websocket_passthrough_request(
|
||||
websocket=websocket,
|
||||
target="wss://aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent",
|
||||
custom_headers={"Authorization": "Bearer token"},
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
forward_headers=False,
|
||||
endpoint="/vertex_ai/live",
|
||||
accept_websocket=False,
|
||||
)
|
||||
|
||||
websocket.close.assert_awaited_once_with()
|
||||
|
||||
|
||||
async def _run_setup_rewrite_passthrough(setup_model: str, llm_router) -> str:
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
_build_vertex_live_setup_model_rewriter,
|
||||
)
|
||||
|
||||
upstream_ws = RecordingUpstreamWebSocket()
|
||||
setup_frame = json.dumps({"setup": {"model": setup_model, "generationConfig": {"responseModalities": ["TEXT"]}}})
|
||||
websocket = _client_websocket(
|
||||
AsyncMock(
|
||||
side_effect=[
|
||||
{"type": "websocket.receive", "text": setup_frame},
|
||||
{"type": "websocket.disconnect"},
|
||||
]
|
||||
)
|
||||
)
|
||||
|
||||
with _patched_websocket_passthrough_environment(upstream_ws):
|
||||
await websocket_passthrough_request(
|
||||
websocket=websocket,
|
||||
target="wss://aiplatform.googleapis.com/ws/google.cloud.aiplatform.v1.LlmBidiService/BidiGenerateContent",
|
||||
custom_headers={"Authorization": "Bearer token"},
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
forward_headers=False,
|
||||
endpoint="/vertex_ai/live",
|
||||
accept_websocket=False,
|
||||
setup_model_rewriter=_build_vertex_live_setup_model_rewriter(
|
||||
vertex_project="proj-db",
|
||||
vertex_location="global",
|
||||
llm_router=llm_router,
|
||||
),
|
||||
)
|
||||
|
||||
upstream_ws.send.assert_awaited_once()
|
||||
return upstream_ws.send.await_args.args[0]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"setup_model",
|
||||
["gemini-live-2.5-flash", "publishers/google/models/gemini-live-2.5-flash"],
|
||||
)
|
||||
async def test_websocket_passthrough_rewrites_setup_model_to_full_resource(setup_model):
|
||||
sent_frame = await _run_setup_rewrite_passthrough(setup_model, llm_router=None)
|
||||
|
||||
sent_setup = json.loads(sent_frame)["setup"]
|
||||
assert sent_setup["model"] == "projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash"
|
||||
assert sent_setup["generationConfig"] == {"responseModalities": ["TEXT"]}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_rewrites_gateway_alias_setup_model():
|
||||
llm_router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gemini-live",
|
||||
"litellm_params": {"model": "vertex_ai/gemini-live-2.5-flash"},
|
||||
}
|
||||
]
|
||||
)
|
||||
|
||||
sent_frame = await _run_setup_rewrite_passthrough("gemini-live", llm_router=llm_router)
|
||||
|
||||
sent_setup = json.loads(sent_frame)["setup"]
|
||||
assert sent_setup["model"] == "projects/proj-db/locations/global/publishers/google/models/gemini-live-2.5-flash"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_websocket_passthrough_leaves_full_resource_setup_model_untouched():
|
||||
full_resource = "projects/other/locations/us-central1/publishers/google/models/gemini-2.0-flash-live-preview-04-09"
|
||||
|
||||
sent_frame = await _run_setup_rewrite_passthrough(full_resource, llm_router=None)
|
||||
|
||||
assert json.loads(sent_frame)["setup"]["model"] == full_resource
|
||||
assert sent_frame == json.dumps(
|
||||
{"setup": {"model": full_resource, "generationConfig": {"responseModalities": ["TEXT"]}}}
|
||||
)
|
||||
|
||||
|
||||
def _passthrough_kwargs_for_reservation(
|
||||
user_api_key_dict: UserAPIKeyAuth, parsed_body: Optional[dict] = None
|
||||
) -> dict:
|
||||
|
|
|
|||
|
|
@ -172,3 +172,144 @@ def test_returns_none_when_no_router_and_no_env():
|
|||
passthrough_router = _passthrough_router(None)
|
||||
|
||||
assert passthrough_router.get_credentials(custom_llm_provider="openai", region_name=None) is None
|
||||
|
||||
|
||||
def _vertex_credential(name: str, values: dict) -> CredentialItem:
|
||||
return CredentialItem(credential_name=name, credential_values=values, credential_info={})
|
||||
|
||||
|
||||
def _vertex_deployment(model_name: str, model: str, **litellm_params) -> dict:
|
||||
return {
|
||||
"model_name": model_name,
|
||||
"litellm_params": {"model": model, "use_in_pass_through": True, **litellm_params},
|
||||
}
|
||||
|
||||
|
||||
def test_vertex_deployment_resolves_via_named_credential():
|
||||
CredentialAccessor.upsert_credentials(
|
||||
[
|
||||
_vertex_credential(
|
||||
"cred_gcp",
|
||||
{
|
||||
"vertex_project": "proj-db",
|
||||
"vertex_location": "global",
|
||||
"vertex_credentials": '{"type": "service_account"}',
|
||||
},
|
||||
)
|
||||
]
|
||||
)
|
||||
llm_router = litellm.Router(
|
||||
model_list=[
|
||||
_vertex_deployment(
|
||||
"gemini-live", "vertex_ai/gemini-live-2.5-flash", litellm_credential_name="cred_gcp"
|
||||
)
|
||||
]
|
||||
)
|
||||
passthrough_router = _passthrough_router(llm_router)
|
||||
|
||||
resolved = passthrough_router.get_vertex_credentials_from_router_deployments(model=None)
|
||||
|
||||
assert resolved is not None
|
||||
assert resolved.vertex_project == "proj-db"
|
||||
assert resolved.vertex_location == "global"
|
||||
assert resolved.vertex_credentials == '{"type": "service_account"}'
|
||||
|
||||
|
||||
def test_vertex_deployment_resolves_from_inline_litellm_params():
|
||||
llm_router = litellm.Router(
|
||||
model_list=[
|
||||
_vertex_deployment(
|
||||
"gemini-live",
|
||||
"vertex_ai/gemini-live-2.5-flash",
|
||||
vertex_project="proj-inline",
|
||||
vertex_location="us-east4",
|
||||
vertex_credentials='{"type": "service_account", "project_id": "proj-inline"}',
|
||||
)
|
||||
]
|
||||
)
|
||||
passthrough_router = _passthrough_router(llm_router)
|
||||
|
||||
resolved = passthrough_router.get_vertex_credentials_from_router_deployments(model=None)
|
||||
|
||||
assert resolved is not None
|
||||
assert resolved.vertex_project == "proj-inline"
|
||||
assert resolved.vertex_location == "us-east4"
|
||||
assert resolved.vertex_credentials == '{"type": "service_account", "project_id": "proj-inline"}'
|
||||
|
||||
|
||||
def _two_vertex_deployments_router() -> litellm.Router:
|
||||
return litellm.Router(
|
||||
model_list=[
|
||||
_vertex_deployment(
|
||||
"gemini-flash",
|
||||
"vertex_ai/gemini-2.5-flash",
|
||||
vertex_project="proj-first",
|
||||
vertex_location="us-central1",
|
||||
),
|
||||
_vertex_deployment(
|
||||
"gemini-live",
|
||||
"vertex_ai/gemini-live-2.5-flash",
|
||||
vertex_project="proj-live",
|
||||
vertex_location="global",
|
||||
),
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def test_vertex_model_hint_prefers_matching_deployment():
|
||||
passthrough_router = _passthrough_router(_two_vertex_deployments_router())
|
||||
|
||||
by_alias = passthrough_router.get_vertex_credentials_from_router_deployments(model="gemini-live")
|
||||
by_upstream_id = passthrough_router.get_vertex_credentials_from_router_deployments(
|
||||
model="gemini-live-2.5-flash"
|
||||
)
|
||||
|
||||
assert by_alias is not None and by_alias.vertex_project == "proj-live"
|
||||
assert by_upstream_id is not None and by_upstream_id.vertex_project == "proj-live"
|
||||
|
||||
|
||||
def test_vertex_unmatched_hint_falls_back_to_first_flagged_deployment():
|
||||
passthrough_router = _passthrough_router(_two_vertex_deployments_router())
|
||||
|
||||
unmatched = passthrough_router.get_vertex_credentials_from_router_deployments(model="unknown-model")
|
||||
no_hint = passthrough_router.get_vertex_credentials_from_router_deployments(model=None)
|
||||
|
||||
assert unmatched is not None and unmatched.vertex_project == "proj-first"
|
||||
assert no_hint is not None and no_hint.vertex_project == "proj-first"
|
||||
|
||||
|
||||
def test_no_flagged_vertex_deployment_returns_none():
|
||||
llm_router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gemini-live",
|
||||
"litellm_params": {
|
||||
"model": "vertex_ai/gemini-live-2.5-flash",
|
||||
"vertex_project": "proj-unflagged",
|
||||
"vertex_location": "global",
|
||||
},
|
||||
},
|
||||
_flagged_deployment("openai/gpt-4o", api_key="sk-flagged"),
|
||||
]
|
||||
)
|
||||
passthrough_router = _passthrough_router(llm_router)
|
||||
|
||||
assert passthrough_router.get_vertex_credentials_from_router_deployments(model=None) is None
|
||||
assert _passthrough_router(None).get_vertex_credentials_from_router_deployments(model=None) is None
|
||||
|
||||
|
||||
def test_vertex_deployment_with_deleted_credential_is_skipped(monkeypatch):
|
||||
CredentialAccessor.upsert_credentials(
|
||||
[_vertex_credential("cred_gone", {"vertex_project": "proj-db", "vertex_location": "global"})]
|
||||
)
|
||||
llm_router = litellm.Router(
|
||||
model_list=[
|
||||
_vertex_deployment(
|
||||
"gemini-live", "vertex_ai/gemini-live-2.5-flash", litellm_credential_name="cred_gone"
|
||||
)
|
||||
]
|
||||
)
|
||||
passthrough_router = _passthrough_router(llm_router)
|
||||
monkeypatch.setattr(litellm, "credential_list", [])
|
||||
|
||||
assert passthrough_router.get_vertex_credentials_from_router_deployments(model=None) is None
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue