diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 26f797cf5b2..8c7ec4e7ee7 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -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 diff --git a/litellm/llms/vertex_ai/interactions_passthrough/__init__.py b/litellm/llms/vertex_ai/interactions_passthrough/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/litellm/llms/vertex_ai/interactions_passthrough/id_codec.py b/litellm/llms/vertex_ai/interactions_passthrough/id_codec.py new file mode 100644 index 00000000000..bfb6727dfab --- /dev/null +++ b/litellm/llms/vertex_ai/interactions_passthrough/id_codec.py @@ -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 diff --git a/litellm/llms/vertex_ai/interactions_passthrough/routing.py b/litellm/llms/vertex_ai/interactions_passthrough/routing.py new file mode 100644 index 00000000000..69e6501cd6d --- /dev/null +++ b/litellm/llms/vertex_ai/interactions_passthrough/routing.py @@ -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 diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index c5ab7f1fc63..43193f71dcf 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -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 diff --git a/tests/test_litellm/llms/vertex_ai/interactions_passthrough/test_id_codec.py b/tests/test_litellm/llms/vertex_ai/interactions_passthrough/test_id_codec.py new file mode 100644 index 00000000000..1a3e27daf0f --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/interactions_passthrough/test_id_codec.py @@ -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 diff --git a/tests/test_litellm/llms/vertex_ai/interactions_passthrough/test_routing.py b/tests/test_litellm/llms/vertex_ai/interactions_passthrough/test_routing.py new file mode 100644 index 00000000000..95381fc8928 --- /dev/null +++ b/tests/test_litellm/llms/vertex_ai/interactions_passthrough/test_routing.py @@ -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" diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py index b83d4742b64..7625a39b6ad 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_ai_common_utils.py @@ -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 diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_interactions_passthrough_routing.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_interactions_passthrough_routing.py new file mode 100644 index 00000000000..20d9878a178 --- /dev/null +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_interactions_passthrough_routing.py @@ -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")