mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
feat(vertex): auto-route interactions passthrough by project and location
Vertex AI generateContent passthrough auto-corrects the GCP project and location from the model id in the URL, but the Interactions API cannot: the model is in the request body on create, and get/delete/cancel carry neither a model nor a body. Async interaction polling therefore breaks unless the caller manually tracks the exact project the interaction was created under, since the stateless passthrough picks credentials purely from the URL's project-location. Add opt-in auto-routing behind general_settings vertex_interactions_passthrough_auto_routing (default off). A stateless base64 codec encodes project, location, and the raw provider id into the interaction id returned to the caller. On create the project and location are resolved from the body's model via the router (with a decoded previous_interaction_id as fallback); on get/delete/cancel the opaque id is decoded to route back to the originating project and the raw id is restored before forwarding. No database, no tenant isolation, no enterprise dependency: the id is opaque and stable so polling reuses it verbatim.
This commit is contained in:
Kris Xia
parent
0700b1e54e
commit
3b29a23856
9 changed files with 1314 additions and 10 deletions
|
|
@ -983,6 +983,17 @@ def get_vertex_model_id_from_url(url: str) -> str | None:
|
|||
return match.group(1) if match else None
|
||||
|
||||
|
||||
def is_vertex_interactions_route(url: str) -> bool:
|
||||
return re.search(r"/interactions(?:/[^/]+)?(?::[^/]+)?(?:\?.*)?$", url) is not None
|
||||
|
||||
|
||||
def get_vertex_interaction_id_from_url(url: str) -> str | None:
|
||||
match: Final = re.search(r"/interactions/([^/?:]+)(?::[^/]+)?(?:\?.*)?$", url)
|
||||
if not match:
|
||||
return None
|
||||
return match.group(1)
|
||||
|
||||
|
||||
def replace_project_and_location_in_route(requested_route: str, vertex_project: str, vertex_location: str) -> str:
|
||||
"""
|
||||
Replace project and location values in the route with the provided values
|
||||
|
|
|
|||
54
litellm/llms/vertex_ai/interactions_passthrough/id_codec.py
Normal file
54
litellm/llms/vertex_ai/interactions_passthrough/id_codec.py
Normal file
|
|
@ -0,0 +1,54 @@
|
|||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
_PREFIX: Final = "litellm_proxy"
|
||||
_DISCRIMINATOR: Final = "vertex_interaction"
|
||||
_HEAD: Final = f"{_PREFIX}:{_DISCRIMINATOR};"
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class VertexInteractionId:
|
||||
project: str
|
||||
location: str
|
||||
raw_id: str
|
||||
|
||||
|
||||
def encode(project: str, location: str, raw_id: str) -> str:
|
||||
plaintext: Final = f"{_HEAD}project,{project};location,{location};raw_id,{raw_id}"
|
||||
return base64.urlsafe_b64encode(plaintext.encode()).decode().rstrip("=")
|
||||
|
||||
|
||||
def decode(value: str) -> VertexInteractionId | None:
|
||||
if not isinstance(value, str) or not value:
|
||||
return None
|
||||
padded: Final = value + "=" * (-len(value) % 4)
|
||||
try:
|
||||
plaintext: Final = base64.urlsafe_b64decode(padded).decode()
|
||||
except (ValueError, UnicodeDecodeError):
|
||||
return None
|
||||
if not plaintext.startswith(_HEAD):
|
||||
return None
|
||||
rest: Final = plaintext[len(_HEAD) :]
|
||||
try:
|
||||
project_part, rest2 = rest.split(";", 1)
|
||||
location_part, raw_id_part = rest2.split(";", 1)
|
||||
except ValueError:
|
||||
return None
|
||||
if not (
|
||||
project_part.startswith("project,")
|
||||
and location_part.startswith("location,")
|
||||
and raw_id_part.startswith("raw_id,")
|
||||
):
|
||||
return None
|
||||
return VertexInteractionId(
|
||||
project=project_part[len("project,") :],
|
||||
location=location_part[len("location,") :],
|
||||
raw_id=raw_id_part[len("raw_id,") :],
|
||||
)
|
||||
|
||||
|
||||
def is_encoded(value: str) -> bool:
|
||||
return decode(value) is not None
|
||||
120
litellm/llms/vertex_ai/interactions_passthrough/routing.py
Normal file
120
litellm/llms/vertex_ai/interactions_passthrough/routing.py
Normal file
|
|
@ -0,0 +1,120 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from collections.abc import Callable
|
||||
from dataclasses import dataclass
|
||||
from typing import Final, TypeAlias
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.llms.vertex_ai.common_utils import get_vertex_interaction_id_from_url
|
||||
from litellm.llms.vertex_ai.interactions_passthrough.id_codec import decode, encode
|
||||
|
||||
InteractionBody: TypeAlias = dict[str, object] # mutable-ok: pass-through request state only accepts dict payloads
|
||||
|
||||
|
||||
class _PassThroughDeploymentRouter(BaseModel):
|
||||
model_config = ConfigDict(arbitrary_types_allowed=True, from_attributes=True)
|
||||
get_available_deployment_for_pass_through: Callable[..., object]
|
||||
|
||||
|
||||
class InteractionCreateBody(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
model: str | None = None
|
||||
previous_interaction_id: str | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class ResolvedRoute:
|
||||
project: str | None
|
||||
location: str | None
|
||||
body: InteractionBody
|
||||
|
||||
|
||||
class _DeploymentLiteLLMParams(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
vertex_project: str | None = None
|
||||
vertex_location: str | None = None
|
||||
|
||||
|
||||
class _Deployment(BaseModel):
|
||||
model_config = ConfigDict(extra="ignore")
|
||||
litellm_params: _DeploymentLiteLLMParams | None = None
|
||||
|
||||
|
||||
def _deployment_project_location(model: str, llm_router: object) -> tuple[str | None, str | None]:
|
||||
try:
|
||||
router: Final = _PassThroughDeploymentRouter.model_validate(llm_router)
|
||||
except ValidationError:
|
||||
return None, None
|
||||
try:
|
||||
deployment: Final = router.get_available_deployment_for_pass_through(model=model)
|
||||
except Exception as error: # noqa: BLE001 - router lookup is best-effort; any failure falls back to URL values
|
||||
verbose_proxy_logger.debug("vertex interactions: deployment lookup failed for model %s: %s", model, error)
|
||||
return None, None
|
||||
try:
|
||||
parsed: Final = _Deployment.model_validate(deployment)
|
||||
except ValidationError:
|
||||
return None, None
|
||||
litellm_params: Final = parsed.litellm_params
|
||||
if litellm_params is None:
|
||||
return None, None
|
||||
return litellm_params.vertex_project, litellm_params.vertex_location
|
||||
|
||||
|
||||
def resolve_create_project_location(
|
||||
body: InteractionBody,
|
||||
url_project: str | None,
|
||||
url_location: str | None,
|
||||
llm_router: object,
|
||||
) -> ResolvedRoute:
|
||||
parsed: Final = InteractionCreateBody.model_validate(body)
|
||||
decoded_prev: Final = decode(parsed.previous_interaction_id) if parsed.previous_interaction_id is not None else None
|
||||
forwarded_body: Final[InteractionBody] = (
|
||||
{**body, "previous_interaction_id": decoded_prev.raw_id} # mutable-ok: downstream state requires a dict
|
||||
if decoded_prev is not None
|
||||
else body
|
||||
)
|
||||
prev_project: Final = decoded_prev.project if decoded_prev is not None else None
|
||||
prev_location: Final = decoded_prev.location if decoded_prev is not None else None
|
||||
model_project, model_location = (
|
||||
_deployment_project_location(parsed.model, llm_router) if parsed.model else (None, None)
|
||||
)
|
||||
project: Final = model_project or prev_project or url_project
|
||||
location: Final = model_location or prev_location or url_location
|
||||
return ResolvedRoute(project=project, location=location, body=forwarded_body)
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class InputRewrite:
|
||||
project: str | None
|
||||
location: str | None
|
||||
endpoint: str
|
||||
|
||||
|
||||
def rewrite_interaction_input(
|
||||
endpoint: str,
|
||||
url_project: str | None,
|
||||
url_location: str | None,
|
||||
) -> InputRewrite:
|
||||
interaction_id: Final = get_vertex_interaction_id_from_url(endpoint)
|
||||
if interaction_id is None:
|
||||
return InputRewrite(project=url_project, location=url_location, endpoint=endpoint)
|
||||
decoded: Final = decode(interaction_id)
|
||||
if decoded is None:
|
||||
return InputRewrite(project=url_project, location=url_location, endpoint=endpoint)
|
||||
new_endpoint: Final = endpoint.replace(interaction_id, decoded.raw_id, 1)
|
||||
return InputRewrite(project=decoded.project, location=decoded.location, endpoint=new_endpoint)
|
||||
|
||||
|
||||
def encode_interaction_response_id(
|
||||
response_body: InteractionBody,
|
||||
project: str | None,
|
||||
location: str | None,
|
||||
) -> InteractionBody:
|
||||
if project is None or location is None:
|
||||
return response_body
|
||||
raw_id: Final = response_body.get("id")
|
||||
if not isinstance(raw_id, str) or not raw_id:
|
||||
return response_body
|
||||
return {**response_body, "id": encode(project, location, raw_id)} # mutable-ok: JSON response requires a dict
|
||||
|
|
@ -1812,6 +1812,82 @@ async def _prepare_vertex_auth_headers(
|
|||
)
|
||||
|
||||
|
||||
def _encode_interaction_response(
|
||||
received_value: object,
|
||||
vertex_project: str | None,
|
||||
vertex_location: str | None,
|
||||
) -> object:
|
||||
from starlette.responses import Response as StarletteResponse
|
||||
|
||||
from litellm.llms.vertex_ai.interactions_passthrough.routing import (
|
||||
encode_interaction_response_id,
|
||||
)
|
||||
|
||||
if not isinstance(received_value, StarletteResponse):
|
||||
return received_value
|
||||
body_bytes: Final = getattr(received_value, "body", None)
|
||||
if not isinstance(body_bytes, (bytes, bytearray)):
|
||||
return received_value
|
||||
try:
|
||||
payload: Final = cast(object, json.loads(bytes(body_bytes))) # cast-ok: json.loads is Any
|
||||
except (ValueError, UnicodeDecodeError):
|
||||
return received_value
|
||||
if not isinstance(payload, dict):
|
||||
return received_value
|
||||
typed_payload: Final = cast("dict[str, object]", payload) # cast-ok: isinstance dict above
|
||||
new_payload: Final = encode_interaction_response_id(typed_payload, vertex_project, vertex_location)
|
||||
if new_payload is payload:
|
||||
return received_value
|
||||
preserved_headers: Final = MappingProxyType(
|
||||
{key: value for key, value in received_value.headers.items() if key.lower() != "content-length"}
|
||||
)
|
||||
return StarletteResponse(
|
||||
content=json.dumps(new_payload),
|
||||
status_code=received_value.status_code,
|
||||
media_type="application/json",
|
||||
headers=preserved_headers,
|
||||
)
|
||||
|
||||
|
||||
async def _resolve_interactions_input_routing(
|
||||
endpoint: str,
|
||||
request: Request,
|
||||
vertex_project: str | None,
|
||||
vertex_location: str | None,
|
||||
llm_router: litellm.Router | None,
|
||||
) -> tuple[str, str | None, str | None]:
|
||||
from litellm.llms.vertex_ai.common_utils import get_vertex_interaction_id_from_url
|
||||
from litellm.llms.vertex_ai.interactions_passthrough.routing import (
|
||||
resolve_create_project_location,
|
||||
rewrite_interaction_input,
|
||||
)
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
|
||||
)
|
||||
|
||||
if get_vertex_interaction_id_from_url(endpoint) is not None:
|
||||
rewrite: Final = rewrite_interaction_input(endpoint, vertex_project, vertex_location)
|
||||
return rewrite.endpoint, rewrite.project, rewrite.location
|
||||
|
||||
if request.method == "POST":
|
||||
try:
|
||||
body: Final = cast(object, await request.json()) # cast-ok: request.json() is Any
|
||||
except Exception: # noqa: BLE001 - unreadable/invalid body falls back to URL values without modifying the body
|
||||
return endpoint, vertex_project, vertex_location
|
||||
if isinstance(body, dict):
|
||||
typed_body: Final = cast("dict[str, object]", body) # cast-ok: isinstance dict above
|
||||
resolved: Final = resolve_create_project_location(
|
||||
body=typed_body,
|
||||
url_project=vertex_project,
|
||||
url_location=vertex_location,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
setattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, resolved.body)
|
||||
return endpoint, resolved.project, resolved.location
|
||||
|
||||
return endpoint, vertex_project, vertex_location
|
||||
|
||||
|
||||
async def _base_vertex_proxy_route(
|
||||
endpoint: str,
|
||||
request: Request,
|
||||
|
|
@ -1840,6 +1916,7 @@ async def _base_vertex_proxy_route(
|
|||
get_vertex_location_from_url,
|
||||
get_vertex_model_id_from_url,
|
||||
get_vertex_project_id_from_url,
|
||||
is_vertex_interactions_route,
|
||||
)
|
||||
from litellm.proxy.proxy_server import llm_router
|
||||
|
||||
|
|
@ -1893,12 +1970,32 @@ async def _base_vertex_proxy_route(
|
|||
vertex_location=vertex_location,
|
||||
)
|
||||
|
||||
from litellm.proxy.proxy_server import general_settings as _general_settings
|
||||
|
||||
_general_settings_typed: Final = cast("dict[str, object]", _general_settings) # cast-ok: untyped config dict
|
||||
interactions_auto_routing: Final = bool(
|
||||
_general_settings_typed.get("vertex_interactions_passthrough_auto_routing", False)
|
||||
) and is_vertex_interactions_route(endpoint)
|
||||
|
||||
routed_endpoint, routed_project, routed_location = (
|
||||
await _resolve_interactions_input_routing(
|
||||
endpoint=endpoint,
|
||||
request=request,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
llm_router=llm_router,
|
||||
)
|
||||
if interactions_auto_routing
|
||||
else (endpoint, vertex_project, vertex_location)
|
||||
)
|
||||
routed_encoded_endpoint: Final = httpx.URL(routed_endpoint).path if interactions_auto_routing else encoded_endpoint
|
||||
|
||||
vertex_credentials: Final = passthrough_endpoint_router.get_vertex_credentials(
|
||||
project_id=vertex_project,
|
||||
location=vertex_location,
|
||||
project_id=routed_project,
|
||||
location=routed_location,
|
||||
)
|
||||
|
||||
base_target_url = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location)
|
||||
base_target_url = get_vertex_pass_through_handler.get_default_base_target_url(routed_location)
|
||||
|
||||
# Prepare authentication headers
|
||||
(
|
||||
|
|
@ -1911,8 +2008,8 @@ async def _base_vertex_proxy_route(
|
|||
request=request,
|
||||
vertex_credentials=vertex_credentials,
|
||||
router_credentials=router_credentials,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
vertex_project=routed_project,
|
||||
vertex_location=routed_location,
|
||||
base_target_url=base_target_url,
|
||||
get_vertex_pass_through_handler=get_vertex_pass_through_handler,
|
||||
)
|
||||
|
|
@ -1920,17 +2017,18 @@ async def _base_vertex_proxy_route(
|
|||
if base_target_url is None:
|
||||
base_target_url = get_vertex_base_url(vertex_location)
|
||||
|
||||
request_route: Final = encoded_endpoint
|
||||
request_route: Final = routed_encoded_endpoint
|
||||
verbose_proxy_logger.debug("request_route %s", request_route)
|
||||
|
||||
# Ensure endpoint starts with '/' for proper URL construction
|
||||
if not encoded_endpoint.startswith("/"):
|
||||
encoded_endpoint = "/" + encoded_endpoint
|
||||
normalized_encoded_endpoint: Final = (
|
||||
routed_encoded_endpoint if routed_encoded_endpoint.startswith("/") else "/" + routed_encoded_endpoint
|
||||
)
|
||||
|
||||
# Construct the full target URL using httpx
|
||||
updated_url: Final = construct_target_url(
|
||||
base_url=base_target_url,
|
||||
requested_route=encoded_endpoint,
|
||||
requested_route=normalized_encoded_endpoint,
|
||||
vertex_location=vertex_location,
|
||||
vertex_project=vertex_project,
|
||||
)
|
||||
|
|
@ -1948,7 +2046,7 @@ async def _base_vertex_proxy_route(
|
|||
|
||||
## CREATE PASS-THROUGH
|
||||
endpoint_func: Final = create_pass_through_route(
|
||||
endpoint=endpoint,
|
||||
endpoint=routed_endpoint,
|
||||
target=target,
|
||||
custom_headers=headers,
|
||||
is_streaming_request=is_streaming_request,
|
||||
|
|
@ -1965,6 +2063,9 @@ async def _base_vertex_proxy_route(
|
|||
e.message = f"No credentials found on proxy for project_name={vertex_project} + location={vertex_location}, check `/model/info` for allowed project + region combinations with `use_in_pass_through: true`. Headers were passed through directly but request failed with error: {e.message}"
|
||||
raise e
|
||||
|
||||
if interactions_auto_routing and not is_streaming_request:
|
||||
return _encode_interaction_response(received_value, vertex_project, vertex_location)
|
||||
|
||||
return received_value
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,83 @@
|
|||
import pytest
|
||||
|
||||
from litellm.llms.vertex_ai.interactions_passthrough.id_codec import (
|
||||
VertexInteractionId,
|
||||
decode,
|
||||
encode,
|
||||
is_encoded,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"project, location, raw_id",
|
||||
[
|
||||
("gemini-0610-462508", "global", "video-43dffcd7-2f1f-4dc1-ac05-a8b885f8822d"),
|
||||
("proj-2", "us-central1", "resp_bGl0ZWxsbTpzb21ldGhpbmc"),
|
||||
("p", "global", "id;with;semicolons"),
|
||||
],
|
||||
)
|
||||
def test_round_trip(project, location, raw_id):
|
||||
encoded = encode(project, location, raw_id)
|
||||
decoded = decode(encoded)
|
||||
assert decoded == VertexInteractionId(project=project, location=location, raw_id=raw_id)
|
||||
|
||||
|
||||
def test_encoding_is_deterministic():
|
||||
a = encode("proj", "global", "video-abc")
|
||||
b = encode("proj", "global", "video-abc")
|
||||
assert a == b
|
||||
|
||||
|
||||
def test_encoded_id_has_no_padding_and_is_urlsafe():
|
||||
encoded = encode("proj", "global", "video-abc")
|
||||
assert "=" not in encoded
|
||||
assert "/" not in encoded and "+" not in encoded
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"value",
|
||||
[
|
||||
"video-43dffcd7-2f1f-4dc1-ac05-a8b885f8822d", # raw vertex id
|
||||
"resp_abc", # raw sync id
|
||||
"", # empty
|
||||
"not base64 at all !!!",
|
||||
],
|
||||
)
|
||||
def test_decode_rejects_non_our_ids(value):
|
||||
assert decode(value) is None
|
||||
assert is_encoded(value) is False
|
||||
|
||||
|
||||
def test_decode_rejects_openai_passthrough_managed_id():
|
||||
import base64
|
||||
|
||||
# OpenAI/Azure passthrough codec uses the "passthrough" discriminator.
|
||||
plaintext = "litellm_proxy:passthrough;provider:openai;unified_id,u1;raw_id,batch_x"
|
||||
other = base64.urlsafe_b64encode(plaintext.encode()).decode().rstrip("=")
|
||||
assert decode(other) is None
|
||||
|
||||
|
||||
def test_decode_rejects_non_string():
|
||||
assert decode(None) is None # type: ignore[arg-type]
|
||||
assert decode(123) is None # type: ignore[arg-type]
|
||||
|
||||
|
||||
def test_is_encoded_true_for_our_ids():
|
||||
assert is_encoded(encode("proj", "global", "video-abc")) is True
|
||||
|
||||
|
||||
def _b64(plaintext: str) -> str:
|
||||
import base64
|
||||
|
||||
return base64.urlsafe_b64encode(plaintext.encode()).decode().rstrip("=")
|
||||
|
||||
|
||||
def test_decode_rejects_correct_head_but_too_few_fields():
|
||||
# Correct discriminator head, but the payload has fewer than the 3 expected
|
||||
# ';'-separated fields, so the unpacking split raises ValueError -> None.
|
||||
assert decode(_b64("litellm_proxy:vertex_interaction;project,p")) is None
|
||||
|
||||
|
||||
def test_decode_rejects_correct_head_but_wrong_field_prefixes():
|
||||
# Three fields, correct head, but the field names are not project/location/raw_id.
|
||||
assert decode(_b64("litellm_proxy:vertex_interaction;proj,p;loc,l;rid,r")) is None
|
||||
|
|
@ -0,0 +1,205 @@
|
|||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.llms.vertex_ai.interactions_passthrough.id_codec import decode, encode
|
||||
from litellm.llms.vertex_ai.interactions_passthrough.routing import (
|
||||
encode_interaction_response_id,
|
||||
resolve_create_project_location,
|
||||
rewrite_interaction_input,
|
||||
)
|
||||
|
||||
|
||||
def _router_with_deployment(vertex_project, vertex_location, model="vertex_ai/gemini-omni-flash-preview"):
|
||||
router = MagicMock()
|
||||
router.get_available_deployment_for_pass_through.return_value = {
|
||||
"litellm_params": {
|
||||
"vertex_project": vertex_project,
|
||||
"vertex_location": vertex_location,
|
||||
"model": model,
|
||||
}
|
||||
}
|
||||
return router
|
||||
|
||||
|
||||
def test_resolves_project_location_from_model():
|
||||
router = _router_with_deployment("real-proj", "global")
|
||||
result = resolve_create_project_location(
|
||||
body={"model": "gemini-omni-flash-preview", "input": []},
|
||||
url_project="PLACEHOLDER",
|
||||
url_location="global",
|
||||
llm_router=router,
|
||||
)
|
||||
assert result.project == "real-proj"
|
||||
assert result.location == "global"
|
||||
router.get_available_deployment_for_pass_through.assert_called_once_with(model="gemini-omni-flash-preview")
|
||||
|
||||
|
||||
def test_no_model_falls_back_to_url():
|
||||
router = MagicMock()
|
||||
result = resolve_create_project_location(
|
||||
body={"input": []},
|
||||
url_project="url-proj",
|
||||
url_location="global",
|
||||
llm_router=router,
|
||||
)
|
||||
assert result.project == "url-proj"
|
||||
assert result.location == "global"
|
||||
router.get_available_deployment_for_pass_through.assert_not_called()
|
||||
|
||||
|
||||
def test_unknown_model_falls_back_to_url():
|
||||
router = MagicMock()
|
||||
router.get_available_deployment_for_pass_through.return_value = None
|
||||
result = resolve_create_project_location(
|
||||
body={"model": "not-configured"},
|
||||
url_project="url-proj",
|
||||
url_location="us-central1",
|
||||
llm_router=router,
|
||||
)
|
||||
assert result.project == "url-proj"
|
||||
assert result.location == "us-central1"
|
||||
|
||||
|
||||
def test_previous_interaction_id_is_decoded_and_used_as_fallback():
|
||||
router = MagicMock()
|
||||
router.get_available_deployment_for_pass_through.return_value = None
|
||||
prev = encode("prev-proj", "global", "video-old")
|
||||
result = resolve_create_project_location(
|
||||
body={"model": "unknown", "previous_interaction_id": prev},
|
||||
url_project="PLACEHOLDER",
|
||||
url_location="global",
|
||||
llm_router=router,
|
||||
)
|
||||
assert result.project == "prev-proj"
|
||||
assert result.location == "global"
|
||||
assert result.body["previous_interaction_id"] == "video-old"
|
||||
|
||||
|
||||
def test_model_resolution_wins_over_previous_interaction_id():
|
||||
router = _router_with_deployment("model-proj", "global")
|
||||
prev = encode("prev-proj", "us-central1", "video-old")
|
||||
result = resolve_create_project_location(
|
||||
body={"model": "gemini-omni-flash-preview", "previous_interaction_id": prev},
|
||||
url_project="PLACEHOLDER",
|
||||
url_location="global",
|
||||
llm_router=router,
|
||||
)
|
||||
assert result.project == "model-proj"
|
||||
assert result.body["previous_interaction_id"] == "video-old"
|
||||
|
||||
|
||||
def test_body_is_not_mutated_in_place():
|
||||
router = MagicMock()
|
||||
router.get_available_deployment_for_pass_through.return_value = None
|
||||
prev = encode("prev-proj", "global", "video-old")
|
||||
original = {"model": "x", "previous_interaction_id": prev}
|
||||
resolve_create_project_location(body=original, url_project="p", url_location="global", llm_router=router)
|
||||
assert original["previous_interaction_id"] == prev
|
||||
|
||||
|
||||
def test_input_rewrite_decodes_and_overrides():
|
||||
opaque = encode("real-proj", "global", "video-abc")
|
||||
endpoint = f"/vertex_ai/v1beta1/projects/PLACEHOLDER/locations/global/interactions/{opaque}"
|
||||
result = rewrite_interaction_input(endpoint, url_project="PLACEHOLDER", url_location="global")
|
||||
assert result.project == "real-proj"
|
||||
assert result.location == "global"
|
||||
assert result.endpoint.endswith("/interactions/video-abc")
|
||||
assert opaque not in result.endpoint
|
||||
|
||||
|
||||
def test_input_rewrite_preserves_cancel_suffix():
|
||||
opaque = encode("real-proj", "us-central1", "video-abc")
|
||||
endpoint = f"/vertex_ai/v1beta1/projects/X/locations/global/interactions/{opaque}:cancel"
|
||||
result = rewrite_interaction_input(endpoint, url_project="X", url_location="global")
|
||||
assert result.project == "real-proj"
|
||||
assert result.location == "us-central1"
|
||||
assert result.endpoint.endswith("/interactions/video-abc:cancel")
|
||||
|
||||
|
||||
def test_input_rewrite_noop_for_raw_id():
|
||||
endpoint = "/vertex_ai/v1beta1/projects/real/locations/global/interactions/video-raw"
|
||||
result = rewrite_interaction_input(endpoint, url_project="real", url_location="global")
|
||||
assert result.project == "real"
|
||||
assert result.location == "global"
|
||||
assert result.endpoint == endpoint
|
||||
|
||||
|
||||
def test_input_rewrite_noop_for_create_url_without_id():
|
||||
# A collection-level create URL has no interaction id, so the endpoint and the
|
||||
# URL project/location are returned untouched.
|
||||
endpoint = "/vertex_ai/v1beta1/projects/real/locations/global/interactions"
|
||||
result = rewrite_interaction_input(endpoint, url_project="real", url_location="global")
|
||||
assert result.project == "real"
|
||||
assert result.location == "global"
|
||||
assert result.endpoint == endpoint
|
||||
|
||||
|
||||
def test_resolve_falls_back_when_router_lacks_passthrough_method():
|
||||
# A router object without get_available_deployment_for_pass_through resolves to URL values.
|
||||
router = object()
|
||||
result = resolve_create_project_location(
|
||||
body={"model": "gemini-omni-flash-preview"},
|
||||
url_project="url-proj",
|
||||
url_location="global",
|
||||
llm_router=router,
|
||||
)
|
||||
assert (result.project, result.location) == ("url-proj", "global")
|
||||
|
||||
|
||||
def test_resolve_falls_back_when_router_raises():
|
||||
router = MagicMock()
|
||||
router.get_available_deployment_for_pass_through.side_effect = RuntimeError("boom")
|
||||
result = resolve_create_project_location(
|
||||
body={"model": "gemini-omni-flash-preview"},
|
||||
url_project="url-proj",
|
||||
url_location="us-central1",
|
||||
llm_router=router,
|
||||
)
|
||||
assert (result.project, result.location) == ("url-proj", "us-central1")
|
||||
|
||||
|
||||
def test_resolve_falls_back_when_deployment_has_no_litellm_params():
|
||||
router = MagicMock()
|
||||
router.get_available_deployment_for_pass_through.return_value = {"model_info": {}}
|
||||
result = resolve_create_project_location(
|
||||
body={"model": "gemini-omni-flash-preview"},
|
||||
url_project="url-proj",
|
||||
url_location="global",
|
||||
llm_router=router,
|
||||
)
|
||||
assert (result.project, result.location) == ("url-proj", "global")
|
||||
|
||||
|
||||
def test_output_encode_rewrites_top_level_id():
|
||||
body = {"id": "video-abc", "status": "in_progress", "object": "interaction"}
|
||||
out = encode_interaction_response_id(body, project="real-proj", location="global")
|
||||
assert out["status"] == "in_progress"
|
||||
decoded = decode(out["id"])
|
||||
assert decoded is not None
|
||||
assert (decoded.project, decoded.location, decoded.raw_id) == ("real-proj", "global", "video-abc")
|
||||
|
||||
|
||||
def test_output_encode_is_stable_round_trip_with_input():
|
||||
opaque = encode("real-proj", "global", "video-abc")
|
||||
endpoint = f"/vertex_ai/v1beta1/projects/PLACEHOLDER/locations/global/interactions/{opaque}"
|
||||
rewritten = rewrite_interaction_input(endpoint, "PLACEHOLDER", "global")
|
||||
body = {"id": "video-abc", "status": "completed"}
|
||||
out = encode_interaction_response_id(body, rewritten.project, rewritten.location)
|
||||
assert out["id"] == opaque
|
||||
|
||||
|
||||
def test_output_encode_noop_when_no_id():
|
||||
body = {"status": "in_progress"}
|
||||
out = encode_interaction_response_id(body, "p", "global")
|
||||
assert out == {"status": "in_progress"}
|
||||
|
||||
|
||||
def test_output_encode_noop_when_project_none():
|
||||
body = {"id": "video-abc"}
|
||||
out = encode_interaction_response_id(body, None, "global")
|
||||
assert out["id"] == "video-abc"
|
||||
|
||||
|
||||
def test_output_encode_does_not_mutate_input():
|
||||
body = {"id": "video-abc", "status": "in_progress"}
|
||||
encode_interaction_response_id(body, "p", "global")
|
||||
assert body["id"] == "video-abc"
|
||||
|
|
@ -13,8 +13,10 @@ sys.path.insert(
|
|||
from litellm.llms.vertex_ai.common_utils import (
|
||||
_get_vertex_url,
|
||||
convert_anyof_null_to_nullable,
|
||||
get_vertex_interaction_id_from_url,
|
||||
get_vertex_location_from_url,
|
||||
get_vertex_project_id_from_url,
|
||||
is_vertex_interactions_route,
|
||||
pop_vertex_request_labels,
|
||||
set_schema_property_ordering,
|
||||
supports_response_json_schema,
|
||||
|
|
@ -1606,3 +1608,64 @@ def test_vertex_text_embedding_request_includes_labels_from_metadata():
|
|||
},
|
||||
)
|
||||
assert req.get("labels") == {"project_id": "cost-center-1"}
|
||||
|
||||
|
||||
class TestVertexInteractionsUrlClassification:
|
||||
def test_is_interactions_route_create(self):
|
||||
url = "/vertex_ai/v1beta1/projects/p/locations/global/interactions"
|
||||
assert is_vertex_interactions_route(url) is True
|
||||
|
||||
def test_is_interactions_route_get(self):
|
||||
url = "/vertex_ai/v1beta1/projects/p/locations/global/interactions/video-abc"
|
||||
assert is_vertex_interactions_route(url) is True
|
||||
|
||||
def test_is_interactions_route_cancel(self):
|
||||
url = "/vertex_ai/v1beta1/projects/p/locations/global/interactions/video-abc:cancel"
|
||||
assert is_vertex_interactions_route(url) is True
|
||||
|
||||
def test_is_interactions_route_false_for_generate_content(self):
|
||||
url = "/vertex_ai/v1/projects/p/locations/us/publishers/google/models/gemini-2.5-flash:generateContent"
|
||||
assert is_vertex_interactions_route(url) is False
|
||||
|
||||
def test_get_interaction_id_from_get_url(self):
|
||||
url = "/vertex_ai/v1beta1/projects/p/locations/global/interactions/video-abc"
|
||||
assert get_vertex_interaction_id_from_url(url) == "video-abc"
|
||||
|
||||
def test_get_interaction_id_from_cancel_url(self):
|
||||
url = "/vertex_ai/v1beta1/projects/p/locations/global/interactions/video-abc:cancel"
|
||||
assert get_vertex_interaction_id_from_url(url) == "video-abc"
|
||||
|
||||
def test_get_interaction_id_none_for_create(self):
|
||||
url = "/vertex_ai/v1beta1/projects/p/locations/global/interactions"
|
||||
assert get_vertex_interaction_id_from_url(url) is None
|
||||
|
||||
def test_is_interactions_route_false_for_project_named_interactions(self):
|
||||
url = "/vertex_ai/v1beta1/projects/interactions/locations/global/publishers/google/models/gemini-2.5-flash:generateContent"
|
||||
assert is_vertex_interactions_route(url) is False
|
||||
|
||||
def test_get_interaction_id_none_for_project_named_interactions(self):
|
||||
url = "/vertex_ai/v1beta1/projects/interactions/locations/global/publishers/google/models/m:generateContent"
|
||||
assert get_vertex_interaction_id_from_url(url) is None
|
||||
|
||||
# Short form (no projects/locations in the URL, like generateContent);
|
||||
# litellm fills project/location in from the resolved deployment.
|
||||
def test_is_interactions_route_short_form_create(self):
|
||||
assert is_vertex_interactions_route("v1beta1/interactions") is True
|
||||
|
||||
def test_is_interactions_route_short_form_get(self):
|
||||
assert is_vertex_interactions_route("v1beta1/interactions/video-abc") is True
|
||||
|
||||
def test_is_interactions_route_short_form_cancel(self):
|
||||
assert is_vertex_interactions_route("v1beta1/interactions/video-abc:cancel") is True
|
||||
|
||||
def test_is_interactions_route_short_form_with_query(self):
|
||||
assert is_vertex_interactions_route("v1beta1/interactions?alt=sse") is True
|
||||
|
||||
def test_get_interaction_id_short_form_get(self):
|
||||
assert get_vertex_interaction_id_from_url("v1beta1/interactions/video-abc") == "video-abc"
|
||||
|
||||
def test_get_interaction_id_short_form_cancel(self):
|
||||
assert get_vertex_interaction_id_from_url("v1beta1/interactions/video-abc:cancel") == "video-abc"
|
||||
|
||||
def test_get_interaction_id_none_for_short_form_create(self):
|
||||
assert get_vertex_interaction_id_from_url("v1beta1/interactions") is None
|
||||
|
|
|
|||
|
|
@ -0,0 +1,667 @@
|
|||
import json
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import Response
|
||||
|
||||
from litellm.llms.vertex_ai.interactions_passthrough.id_codec import decode, encode
|
||||
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
|
||||
_base_vertex_proxy_route,
|
||||
_encode_interaction_response,
|
||||
_resolve_interactions_input_routing,
|
||||
)
|
||||
from litellm.types.passthrough_endpoints.pass_through_endpoints import (
|
||||
LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY,
|
||||
)
|
||||
|
||||
|
||||
def _mock_request(method: str, body: dict | None):
|
||||
request = MagicMock()
|
||||
request.method = method
|
||||
request.headers = {}
|
||||
state = MagicMock()
|
||||
request.state = state
|
||||
request.json = AsyncMock(return_value=(body or {}))
|
||||
return request, state
|
||||
|
||||
|
||||
def _mock_router(vertex_project, vertex_location):
|
||||
router = MagicMock()
|
||||
router.get_available_deployment_for_pass_through.return_value = {
|
||||
"litellm_params": {
|
||||
"vertex_project": vertex_project,
|
||||
"vertex_location": vertex_location,
|
||||
"model": "vertex_ai/gemini-omni-flash-preview",
|
||||
}
|
||||
}
|
||||
return router
|
||||
|
||||
|
||||
def _make_response(payload: dict) -> Response:
|
||||
return Response(content=json.dumps(payload), status_code=200, media_type="application/json")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_auto_routes_and_encodes_id():
|
||||
request, state = _mock_request("POST", {"model": "gemini-omni-flash-preview", "background": True})
|
||||
fastapi_response = MagicMock()
|
||||
handler = MagicMock()
|
||||
handler.get_default_base_target_url.return_value = "https://aiplatform.googleapis.com"
|
||||
router = _mock_router("real-proj", "global")
|
||||
|
||||
captured = {}
|
||||
|
||||
def fake_create_pass_through_route(endpoint, target, custom_headers, is_streaming_request):
|
||||
captured["target"] = target
|
||||
|
||||
async def endpoint_func(request, fastapi_response, user_api_key_dict):
|
||||
return _make_response({"id": "video-abc", "status": "in_progress"})
|
||||
|
||||
return endpoint_func
|
||||
|
||||
async def _echo_prep_headers(**kwargs):
|
||||
return (
|
||||
{},
|
||||
"https://aiplatform.googleapis.com",
|
||||
False,
|
||||
kwargs["vertex_project"],
|
||||
kwargs["vertex_location"],
|
||||
)
|
||||
|
||||
endpoint = "/vertex_ai/v1beta1/projects/PLACEHOLDER/locations/global/interactions"
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
|
||||
new=AsyncMock(return_value=MagicMock()),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router"
|
||||
) as mock_pt_router,
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route",
|
||||
side_effect=fake_create_pass_through_route,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers",
|
||||
new=AsyncMock(side_effect=_echo_prep_headers),
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.llm_router", new=router),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
new={"vertex_interactions_passthrough_auto_routing": True},
|
||||
),
|
||||
):
|
||||
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
|
||||
result = await _base_vertex_proxy_route(
|
||||
endpoint=endpoint,
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
get_vertex_pass_through_handler=handler,
|
||||
)
|
||||
|
||||
assert "projects/real-proj/" in captured["target"]
|
||||
payload = json.loads(bytes(result.body))
|
||||
decoded = decode(payload["id"])
|
||||
assert decoded is not None
|
||||
assert (decoded.project, decoded.location, decoded.raw_id) == ("real-proj", "global", "video-abc")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_short_form_url_without_project_auto_routes():
|
||||
# Short form like generateContent: no projects/locations in the URL at all.
|
||||
request, state = _mock_request("POST", {"model": "gemini-omni-flash-preview", "background": True})
|
||||
fastapi_response = MagicMock()
|
||||
handler = MagicMock()
|
||||
handler.get_default_base_target_url.return_value = "https://aiplatform.googleapis.com"
|
||||
router = _mock_router("real-proj", "global")
|
||||
|
||||
captured = {}
|
||||
|
||||
def fake_create_pass_through_route(endpoint, target, custom_headers, is_streaming_request):
|
||||
captured["target"] = target
|
||||
|
||||
async def endpoint_func(request, fastapi_response, user_api_key_dict):
|
||||
return _make_response({"id": "video-abc", "status": "in_progress"})
|
||||
|
||||
return endpoint_func
|
||||
|
||||
async def _echo_prep_headers(**kwargs):
|
||||
return (
|
||||
{},
|
||||
"https://aiplatform.googleapis.com",
|
||||
False,
|
||||
kwargs["vertex_project"],
|
||||
kwargs["vertex_location"],
|
||||
)
|
||||
|
||||
endpoint = "v1beta1/interactions"
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
|
||||
new=AsyncMock(return_value=MagicMock()),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router"
|
||||
) as mock_pt_router,
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route",
|
||||
side_effect=fake_create_pass_through_route,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers",
|
||||
new=AsyncMock(side_effect=_echo_prep_headers),
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.llm_router", new=router),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
new={"vertex_interactions_passthrough_auto_routing": True},
|
||||
),
|
||||
):
|
||||
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
|
||||
result = await _base_vertex_proxy_route(
|
||||
endpoint=endpoint,
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
get_vertex_pass_through_handler=handler,
|
||||
)
|
||||
|
||||
# litellm filled in the resolved project/location for the projectless URL.
|
||||
assert "projects/real-proj/locations/global/interactions" in captured["target"]
|
||||
payload = json.loads(bytes(result.body))
|
||||
decoded = decode(payload["id"])
|
||||
assert decoded is not None
|
||||
assert (decoded.project, decoded.location, decoded.raw_id) == ("real-proj", "global", "video-abc")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_short_form_url_routes_back_via_opaque_id():
|
||||
# Short form get: no projects/locations; the opaque id supplies them.
|
||||
opaque = encode("real-proj", "global", "video-abc")
|
||||
request, state = _mock_request("GET", None)
|
||||
fastapi_response = MagicMock()
|
||||
handler = MagicMock()
|
||||
handler.get_default_base_target_url.return_value = "https://aiplatform.googleapis.com"
|
||||
|
||||
captured = {}
|
||||
|
||||
def fake_create_pass_through_route(endpoint, target, custom_headers, is_streaming_request):
|
||||
captured["target"] = target
|
||||
|
||||
async def endpoint_func(request, fastapi_response, user_api_key_dict):
|
||||
return _make_response({"id": "video-abc", "status": "completed"})
|
||||
|
||||
return endpoint_func
|
||||
|
||||
async def _echo_prep_headers(**kwargs):
|
||||
return (
|
||||
{},
|
||||
"https://aiplatform.googleapis.com",
|
||||
False,
|
||||
kwargs["vertex_project"],
|
||||
kwargs["vertex_location"],
|
||||
)
|
||||
|
||||
endpoint = f"v1beta1/interactions/{opaque}"
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
|
||||
new=AsyncMock(return_value=MagicMock()),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router"
|
||||
) as mock_pt_router,
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route",
|
||||
side_effect=fake_create_pass_through_route,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers",
|
||||
new=AsyncMock(side_effect=_echo_prep_headers),
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.llm_router", new=MagicMock()),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
new={"vertex_interactions_passthrough_auto_routing": True},
|
||||
),
|
||||
):
|
||||
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
|
||||
result = await _base_vertex_proxy_route(
|
||||
endpoint=endpoint,
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
get_vertex_pass_through_handler=handler,
|
||||
)
|
||||
|
||||
# Routed to the decoded project with the raw id, project filled in by litellm.
|
||||
assert "projects/real-proj/locations/global/interactions/video-abc" in captured["target"]
|
||||
payload = json.loads(bytes(result.body))
|
||||
assert payload["id"] == opaque
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_routes_back_via_opaque_id_ignoring_url_project():
|
||||
opaque = encode("real-proj", "global", "video-abc")
|
||||
request, state = _mock_request("GET", None)
|
||||
fastapi_response = MagicMock()
|
||||
handler = MagicMock()
|
||||
handler.get_default_base_target_url.return_value = "https://aiplatform.googleapis.com"
|
||||
|
||||
captured = {}
|
||||
|
||||
def fake_create_pass_through_route(endpoint, target, custom_headers, is_streaming_request):
|
||||
captured["target"] = target
|
||||
|
||||
async def endpoint_func(request, fastapi_response, user_api_key_dict):
|
||||
return _make_response({"id": "video-abc", "status": "completed"})
|
||||
|
||||
return endpoint_func
|
||||
|
||||
async def _echo_prep_headers(**kwargs):
|
||||
return (
|
||||
{},
|
||||
"https://aiplatform.googleapis.com",
|
||||
False,
|
||||
kwargs["vertex_project"],
|
||||
kwargs["vertex_location"],
|
||||
)
|
||||
|
||||
endpoint = f"/vertex_ai/v1beta1/projects/PLACEHOLDER/locations/global/interactions/{opaque}"
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
|
||||
new=AsyncMock(return_value=MagicMock()),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router"
|
||||
) as mock_pt_router,
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route",
|
||||
side_effect=fake_create_pass_through_route,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers",
|
||||
new=AsyncMock(side_effect=_echo_prep_headers),
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.llm_router", new=MagicMock()),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
new={"vertex_interactions_passthrough_auto_routing": True},
|
||||
),
|
||||
):
|
||||
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
|
||||
result = await _base_vertex_proxy_route(
|
||||
endpoint=endpoint,
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
get_vertex_pass_through_handler=handler,
|
||||
)
|
||||
|
||||
assert "projects/real-proj/" in captured["target"]
|
||||
assert "interactions/video-abc" in captured["target"]
|
||||
payload = json.loads(bytes(result.body))
|
||||
assert payload["id"] == opaque
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_flag_off_leaves_url_untouched():
|
||||
request, state = _mock_request("POST", {"model": "gemini-omni-flash-preview"})
|
||||
fastapi_response = MagicMock()
|
||||
handler = MagicMock()
|
||||
handler.get_default_base_target_url.return_value = "https://aiplatform.googleapis.com"
|
||||
router = _mock_router("real-proj", "global")
|
||||
|
||||
captured = {}
|
||||
|
||||
def fake_create_pass_through_route(endpoint, target, custom_headers, is_streaming_request):
|
||||
captured["target"] = target
|
||||
|
||||
async def endpoint_func(request, fastapi_response, user_api_key_dict):
|
||||
return _make_response({"id": "video-abc", "status": "in_progress"})
|
||||
|
||||
return endpoint_func
|
||||
|
||||
async def _echo_prep_headers(**kwargs):
|
||||
return (
|
||||
{},
|
||||
"https://aiplatform.googleapis.com",
|
||||
False,
|
||||
kwargs["vertex_project"],
|
||||
kwargs["vertex_location"],
|
||||
)
|
||||
|
||||
endpoint = "/vertex_ai/v1beta1/projects/PLACEHOLDER/locations/global/interactions"
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
|
||||
new=AsyncMock(return_value=MagicMock()),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router"
|
||||
) as mock_pt_router,
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route",
|
||||
side_effect=fake_create_pass_through_route,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers",
|
||||
new=AsyncMock(side_effect=_echo_prep_headers),
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.llm_router", new=router),
|
||||
patch("litellm.proxy.proxy_server.general_settings", new={}),
|
||||
):
|
||||
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
|
||||
result = await _base_vertex_proxy_route(
|
||||
endpoint=endpoint,
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
get_vertex_pass_through_handler=handler,
|
||||
)
|
||||
|
||||
assert "projects/PLACEHOLDER/" in captured["target"]
|
||||
payload = json.loads(bytes(result.body))
|
||||
assert payload["id"] == "video-abc"
|
||||
assert decode(payload["id"]) is None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_response_preserves_upstream_headers():
|
||||
request, state = _mock_request("POST", {"model": "gemini-omni-flash-preview"})
|
||||
fastapi_response = MagicMock()
|
||||
handler = MagicMock()
|
||||
handler.get_default_base_target_url.return_value = "https://aiplatform.googleapis.com"
|
||||
router = _mock_router("real-proj", "global")
|
||||
|
||||
def fake_create_pass_through_route(endpoint, target, custom_headers, is_streaming_request):
|
||||
async def endpoint_func(request, fastapi_response, user_api_key_dict):
|
||||
return Response(
|
||||
content=json.dumps({"id": "video-abc", "status": "in_progress"}),
|
||||
status_code=200,
|
||||
media_type="application/json",
|
||||
headers={"x-goog-request-id": "trace-123"},
|
||||
)
|
||||
|
||||
return endpoint_func
|
||||
|
||||
async def _echo_prep_headers(**kwargs):
|
||||
return (
|
||||
{},
|
||||
"https://aiplatform.googleapis.com",
|
||||
False,
|
||||
kwargs["vertex_project"],
|
||||
kwargs["vertex_location"],
|
||||
)
|
||||
|
||||
endpoint = "/vertex_ai/v1beta1/projects/PLACEHOLDER/locations/global/interactions"
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
|
||||
new=AsyncMock(return_value=MagicMock()),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router"
|
||||
) as mock_pt_router,
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route",
|
||||
side_effect=fake_create_pass_through_route,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers",
|
||||
new=AsyncMock(side_effect=_echo_prep_headers),
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.llm_router", new=router),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
new={"vertex_interactions_passthrough_auto_routing": True},
|
||||
),
|
||||
):
|
||||
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
|
||||
result = await _base_vertex_proxy_route(
|
||||
endpoint=endpoint,
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
get_vertex_pass_through_handler=handler,
|
||||
)
|
||||
|
||||
assert result.headers["x-goog-request-id"] == "trace-123"
|
||||
assert int(result.headers["content-length"]) == len(bytes(result.body))
|
||||
payload = json.loads(bytes(result.body))
|
||||
assert decode(payload["id"]) is not None
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_round_trip_through_real_prepare_auth_headers():
|
||||
from litellm.types.passthrough_endpoints.vertex_ai import (
|
||||
VertexPassThroughCredentials,
|
||||
)
|
||||
|
||||
opaque = encode("real-proj", "global", "video-abc")
|
||||
request, state = _mock_request("GET", None)
|
||||
fastapi_response = MagicMock()
|
||||
handler = MagicMock()
|
||||
handler.get_default_base_target_url.return_value = "https://aiplatform.googleapis.com"
|
||||
handler.update_base_target_url_with_credential_location.side_effect = lambda base_url, location: base_url
|
||||
|
||||
captured = {}
|
||||
|
||||
def fake_create_pass_through_route(endpoint, target, custom_headers, is_streaming_request):
|
||||
captured["target"] = target
|
||||
|
||||
async def endpoint_func(request, fastapi_response, user_api_key_dict):
|
||||
return _make_response({"id": "video-abc", "status": "completed"})
|
||||
|
||||
return endpoint_func
|
||||
|
||||
endpoint = f"/vertex_ai/v1beta1/projects/PLACEHOLDER/locations/global/interactions/{opaque}"
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
|
||||
new=AsyncMock(return_value=MagicMock()),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router"
|
||||
) as mock_pt_router,
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route",
|
||||
side_effect=fake_create_pass_through_route,
|
||||
),
|
||||
patch(
|
||||
"litellm.llms.vertex_ai.vertex_llm_base.VertexBase._ensure_access_token_async",
|
||||
new=AsyncMock(side_effect=lambda credentials, project_id, custom_llm_provider: ("tok", project_id)),
|
||||
),
|
||||
patch(
|
||||
"litellm.llms.vertex_ai.vertex_llm_base.VertexBase._get_token_and_url",
|
||||
new=MagicMock(return_value=("tok", "https://aiplatform.googleapis.com")),
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.llm_router", new=MagicMock()),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
new={"vertex_interactions_passthrough_auto_routing": True},
|
||||
),
|
||||
):
|
||||
# A credential registered for the decoded project: its vertex_project
|
||||
# equals the lookup key, so _prepare_vertex_auth_headers is idempotent.
|
||||
mock_pt_router.get_vertex_credentials.return_value = VertexPassThroughCredentials(
|
||||
vertex_project="real-proj",
|
||||
vertex_location="global",
|
||||
vertex_credentials="/fake/creds.json",
|
||||
)
|
||||
result = await _base_vertex_proxy_route(
|
||||
endpoint=endpoint,
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
get_vertex_pass_through_handler=handler,
|
||||
)
|
||||
|
||||
# Credentials were looked up by the DECODED project, not the placeholder URL.
|
||||
assert mock_pt_router.get_vertex_credentials.call_args.kwargs["project_id"] == "real-proj"
|
||||
# Routed to the decoded project with the raw id, through the real auth-header path.
|
||||
assert "projects/real-proj/" in captured["target"]
|
||||
assert "interactions/video-abc" in captured["target"]
|
||||
# Response re-encodes to the SAME opaque id the caller sent (stable polling).
|
||||
payload = json.loads(bytes(result.body))
|
||||
assert payload["id"] == opaque
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_streaming_create_response_passes_through_untouched():
|
||||
from fastapi.responses import StreamingResponse
|
||||
|
||||
request, state = _mock_request("POST", {"model": "gemini-omni-flash-preview", "stream": True})
|
||||
fastapi_response = MagicMock()
|
||||
handler = MagicMock()
|
||||
handler.get_default_base_target_url.return_value = "https://aiplatform.googleapis.com"
|
||||
router = _mock_router("real-proj", "global")
|
||||
|
||||
captured = {}
|
||||
sentinel = StreamingResponse(iter([b"data: {}\n\n"]), media_type="text/event-stream")
|
||||
|
||||
def fake_create_pass_through_route(endpoint, target, custom_headers, is_streaming_request):
|
||||
captured["target"] = target
|
||||
|
||||
async def endpoint_func(request, fastapi_response, user_api_key_dict):
|
||||
return sentinel
|
||||
|
||||
return endpoint_func
|
||||
|
||||
async def _echo_prep_headers(**kwargs):
|
||||
return (
|
||||
{},
|
||||
"https://aiplatform.googleapis.com",
|
||||
False,
|
||||
kwargs["vertex_project"],
|
||||
kwargs["vertex_location"],
|
||||
)
|
||||
|
||||
endpoint = "/vertex_ai/v1beta1/projects/PLACEHOLDER/locations/global/interactions"
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
|
||||
new=AsyncMock(return_value=MagicMock()),
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router"
|
||||
) as mock_pt_router,
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route",
|
||||
side_effect=fake_create_pass_through_route,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers",
|
||||
new=AsyncMock(side_effect=_echo_prep_headers),
|
||||
),
|
||||
patch("litellm.proxy.proxy_server.llm_router", new=router),
|
||||
patch(
|
||||
"litellm.proxy.proxy_server.general_settings",
|
||||
new={"vertex_interactions_passthrough_auto_routing": True},
|
||||
),
|
||||
):
|
||||
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
|
||||
result = await _base_vertex_proxy_route(
|
||||
endpoint=endpoint,
|
||||
request=request,
|
||||
fastapi_response=fastapi_response,
|
||||
get_vertex_pass_through_handler=handler,
|
||||
)
|
||||
|
||||
# Input-side model resolution still applied for a streaming create.
|
||||
assert "projects/real-proj/" in captured["target"]
|
||||
# A StreamingResponse has no `.body`, so the output encode degrades to a
|
||||
# no-op and the streaming response is returned untouched (never buffered).
|
||||
assert result is sentinel
|
||||
|
||||
|
||||
def test_encode_response_noop_for_non_response_return():
|
||||
obj = object()
|
||||
assert _encode_interaction_response(obj, "proj", "global") is obj
|
||||
|
||||
|
||||
def test_encode_response_noop_for_non_json_body():
|
||||
resp = Response(content=b"not json", status_code=200, media_type="text/plain")
|
||||
assert _encode_interaction_response(resp, "proj", "global") is resp
|
||||
|
||||
|
||||
def test_encode_response_noop_for_non_dict_json_body():
|
||||
resp = Response(content=json.dumps([1, 2, 3]), status_code=200, media_type="application/json")
|
||||
assert _encode_interaction_response(resp, "proj", "global") is resp
|
||||
|
||||
|
||||
def test_encode_response_noop_when_no_id_field():
|
||||
resp = Response(content=json.dumps({"status": "in_progress"}), status_code=200, media_type="application/json")
|
||||
assert _encode_interaction_response(resp, "proj", "global") is resp
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_input_routing_body_read_error_falls_back_to_url():
|
||||
request = MagicMock()
|
||||
request.method = "POST"
|
||||
request.state = SimpleNamespace()
|
||||
request.json = AsyncMock(side_effect=ValueError("bad body"))
|
||||
endpoint = "/vertex_ai/v1beta1/projects/url-proj/locations/global/interactions"
|
||||
|
||||
new_endpoint, project, location = await _resolve_interactions_input_routing(
|
||||
endpoint=endpoint,
|
||||
request=request,
|
||||
vertex_project="url-proj",
|
||||
vertex_location="global",
|
||||
llm_router=MagicMock(),
|
||||
)
|
||||
assert new_endpoint == endpoint
|
||||
assert (project, location) == ("url-proj", "global")
|
||||
assert not hasattr(request.state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resolve_input_routing_non_post_non_id_url_is_untouched():
|
||||
request = MagicMock()
|
||||
request.method = "DELETE"
|
||||
request.state = MagicMock()
|
||||
# A DELETE on the collection-level URL (no id): neither branch applies, values pass through.
|
||||
endpoint = "/vertex_ai/v1beta1/projects/url-proj/locations/global/interactions"
|
||||
|
||||
new_endpoint, project, location = await _resolve_interactions_input_routing(
|
||||
endpoint=endpoint,
|
||||
request=request,
|
||||
vertex_project="url-proj",
|
||||
vertex_location="global",
|
||||
llm_router=MagicMock(),
|
||||
)
|
||||
assert new_endpoint == endpoint
|
||||
assert (project, location) == ("url-proj", "global")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_forwards_decoded_previous_interaction_id_upstream():
|
||||
# A caller passes back an opaque previous_interaction_id it received earlier.
|
||||
# The body forwarded upstream must carry the DECODED raw id, not the opaque
|
||||
# string, or Vertex cannot parse it. The rewritten body is stashed on
|
||||
# request.state under LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY, which is what
|
||||
# pass_through_request forwards.
|
||||
prev_opaque = encode("real-proj", "global", "video-prev")
|
||||
request, state = _mock_request(
|
||||
"POST",
|
||||
{"model": "gemini-omni-flash-preview", "previous_interaction_id": prev_opaque},
|
||||
)
|
||||
|
||||
new_endpoint, project, location = await _resolve_interactions_input_routing(
|
||||
endpoint="v1beta1/interactions",
|
||||
request=request,
|
||||
vertex_project=None,
|
||||
vertex_location=None,
|
||||
llm_router=_mock_router("real-proj", "global"),
|
||||
)
|
||||
|
||||
# The body handed to pass_through_request carries the raw id, not the opaque one.
|
||||
forwarded_body = getattr(state, LITELLM_PASS_THROUGH_CUSTOM_BODY_STATE_KEY)
|
||||
assert forwarded_body["previous_interaction_id"] == "video-prev"
|
||||
assert forwarded_body["previous_interaction_id"] != prev_opaque
|
||||
# Model resolution still drives project/location on create.
|
||||
assert (project, location) == ("real-proj", "global")
|
||||
Loading…
Add table
Reference in a new issue