mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
fix(interactions): route Vertex lifecycle calls through the selected deployment (#51511)
* fix(interactions): route Vertex interaction follow-ups to the selected deployment GET, cancel and DELETE by interaction id accept a model group or a deployment id through the model query parameter or the x-litellm-model header, authorize it against the key and team like any model, and run with that deployment's project and credentials. Without a selector they still go to Google AI Studio. Vertex cancel now uses the documented /cancel path. * test(bedrock): pin the SigV4 clock in the embedding signature test SigV4Auth.add_auth stamps the request with botocore.auth.get_current_datetime(), overwriting the timestamp the test set on the replay. When the proxy call and the replay landed in different seconds, the signatures differed. Fix the module clock for both signings and drop the ineffective manual timestamp. --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
This commit is contained in:
joshua-berri
GitHub
parent
c5843a8124
commit
ac5fd9e0d0
7 changed files with 554 additions and 24 deletions
|
|
@ -146,4 +146,4 @@ class VertexAIInteractionsConfig(VertexBase, GoogleAIStudioInteractionsConfig):
|
|||
litellm_params: GenericLiteLLMParams,
|
||||
headers: Mapping[str, str],
|
||||
) -> tuple[str, dict]: # mutable-ok: BaseInteractionsAPIConfig declares a plain-dict request body
|
||||
return self._interaction_by_id_request(interaction_id, api_base, litellm_params, url_suffix=":cancel")
|
||||
return self._interaction_by_id_request(interaction_id, api_base, litellm_params, url_suffix="/cancel")
|
||||
|
|
|
|||
|
|
@ -2423,12 +2423,14 @@ _MODEL_ROUTING_ROUTE_MARKERS: Final = (
|
|||
"/evals",
|
||||
"/fine_tuning",
|
||||
"/videos",
|
||||
"/interactions/",
|
||||
)
|
||||
_MODEL_ROUTING_HEADER_OR_QUERY_ROUTE_MARKERS: Final = (
|
||||
"/files",
|
||||
"/batches",
|
||||
"/skills",
|
||||
"/evals",
|
||||
"/interactions/",
|
||||
)
|
||||
_MODEL_ROUTING_QUERY_TARGET_MODEL_ROUTE_MARKERS: Final = (
|
||||
"/files",
|
||||
|
|
@ -2602,6 +2604,19 @@ def video_access_group(
|
|||
return creating_group
|
||||
|
||||
|
||||
_INTERACTION_BY_ID_ROUTE_MARKERS: Final = ("/interactions/",)
|
||||
|
||||
|
||||
def _group_of_deployment_selector(selector: object, llm_router: Router | None) -> object:
|
||||
if not isinstance(selector, str) or llm_router is None:
|
||||
return selector
|
||||
registry: Final = llm_router.registry()
|
||||
if selector in registry.index("model_name"):
|
||||
return selector
|
||||
deployment: Final = registry.deployment(selector)
|
||||
return selector if deployment is None else deployment.model_name
|
||||
|
||||
|
||||
def _resolve_model_id_with_router(model_id: str | None, llm_router: Router | None) -> str | None:
|
||||
if model_id is None or llm_router is None:
|
||||
return model_id
|
||||
|
|
@ -2680,9 +2695,15 @@ def _extract_model_candidates_from_request(
|
|||
uses_completion_model_sources: Final = _route_matches_any_marker(
|
||||
route=route, markers=_MODEL_ROUTING_COMPLETION_MODEL_ROUTE_MARKERS
|
||||
)
|
||||
selects_by_interaction_id: Final = _route_matches_any_marker(route=route, markers=_INTERACTION_BY_ID_ROUTE_MARKERS)
|
||||
|
||||
body_model: Final = request_data.get("model")
|
||||
_append_model_candidates(candidates, body_model)
|
||||
_append_model_candidates(
|
||||
candidates,
|
||||
_group_of_deployment_selector(body_model, llm_router)
|
||||
if selects_by_interaction_id and isinstance(body_model, str)
|
||||
else body_model,
|
||||
)
|
||||
if uses_body_target_model_sources or not body_model:
|
||||
_append_model_candidates(candidates, request_data.get("target_model_names"))
|
||||
if _route_matches_any_marker(route=route, markers=_MODEL_ROUTING_SESSION_MODEL_ROUTE_MARKERS):
|
||||
|
|
@ -2694,14 +2715,14 @@ def _extract_model_candidates_from_request(
|
|||
|
||||
if uses_model_routing_sources:
|
||||
if uses_header_or_query_model_sources:
|
||||
_append_model_candidates(
|
||||
candidates,
|
||||
for selector in (
|
||||
_get_case_insensitive_mapping_value(request_query_params, "model"),
|
||||
)
|
||||
_append_model_candidates(
|
||||
candidates,
|
||||
_get_case_insensitive_mapping_value(request_headers, MODEL_ROUTING_HEADER_NAME),
|
||||
)
|
||||
):
|
||||
_append_model_candidates(
|
||||
candidates,
|
||||
_group_of_deployment_selector(selector, llm_router) if selects_by_interaction_id else selector,
|
||||
)
|
||||
if uses_query_target_model_sources:
|
||||
_append_model_candidates(
|
||||
candidates,
|
||||
|
|
|
|||
|
|
@ -10,6 +10,7 @@ from pydantic import BaseModel, BeforeValidator, ConfigDict, Field, JsonValue, T
|
|||
from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types
|
||||
from litellm.proxy._types import *
|
||||
from litellm.proxy._types import ProxyException, TokenCountRequest
|
||||
from litellm.proxy.auth.auth_utils import MODEL_ROUTING_HEADER_NAME
|
||||
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth
|
||||
from litellm.proxy.common_request_processing import (
|
||||
ProxyBaseLLMRequestProcessing,
|
||||
|
|
@ -487,6 +488,10 @@ async def create_interaction(
|
|||
)
|
||||
|
||||
|
||||
def _interaction_model_group(request: Request) -> str | None:
|
||||
return request.query_params.get("model") or request.headers.get(MODEL_ROUTING_HEADER_NAME)
|
||||
|
||||
|
||||
@router.get(
|
||||
"/v1beta/interactions/{interaction_id}",
|
||||
dependencies=[Depends(user_api_key_auth)],
|
||||
|
|
@ -524,7 +529,8 @@ async def get_interaction(
|
|||
version,
|
||||
)
|
||||
|
||||
data: Final = {"interaction_id": interaction_id, "custom_llm_provider": "gemini"}
|
||||
model: Final = _interaction_model_group(request)
|
||||
data: Final = {"interaction_id": interaction_id, "custom_llm_provider": "gemini", "model": model}
|
||||
|
||||
processor: Final = ProxyBaseLLMRequestProcessing(data=data)
|
||||
try:
|
||||
|
|
@ -538,7 +544,7 @@ async def get_interaction(
|
|||
general_settings=general_settings,
|
||||
proxy_config=proxy_config,
|
||||
select_data_generator=select_data_generator,
|
||||
model=None,
|
||||
model=model,
|
||||
user_model=user_model,
|
||||
user_temperature=user_temperature,
|
||||
user_request_timeout=user_request_timeout,
|
||||
|
|
@ -592,7 +598,8 @@ async def delete_interaction(
|
|||
version,
|
||||
)
|
||||
|
||||
data: Final = {"interaction_id": interaction_id, "custom_llm_provider": "gemini"}
|
||||
model: Final = _interaction_model_group(request)
|
||||
data: Final = {"interaction_id": interaction_id, "custom_llm_provider": "gemini", "model": model}
|
||||
|
||||
processor: Final = ProxyBaseLLMRequestProcessing(data=data)
|
||||
try:
|
||||
|
|
@ -606,7 +613,7 @@ async def delete_interaction(
|
|||
general_settings=general_settings,
|
||||
proxy_config=proxy_config,
|
||||
select_data_generator=select_data_generator,
|
||||
model=None,
|
||||
model=model,
|
||||
user_model=user_model,
|
||||
user_temperature=user_temperature,
|
||||
user_request_timeout=user_request_timeout,
|
||||
|
|
@ -660,7 +667,8 @@ async def cancel_interaction(
|
|||
version,
|
||||
)
|
||||
|
||||
data: Final = {"interaction_id": interaction_id, "custom_llm_provider": "gemini"}
|
||||
model: Final = _interaction_model_group(request)
|
||||
data: Final = {"interaction_id": interaction_id, "custom_llm_provider": "gemini", "model": model}
|
||||
|
||||
processor: Final = ProxyBaseLLMRequestProcessing(data=data)
|
||||
try:
|
||||
|
|
@ -674,7 +682,7 @@ async def cancel_interaction(
|
|||
general_settings=general_settings,
|
||||
proxy_config=proxy_config,
|
||||
select_data_generator=select_data_generator,
|
||||
model=None,
|
||||
model=model,
|
||||
user_model=user_model,
|
||||
user_temperature=user_temperature,
|
||||
user_request_timeout=user_request_timeout,
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import json
|
||||
import asyncio
|
||||
import os
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
from unittest.mock import Mock, patch
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
|
@ -2247,11 +2248,15 @@ def test_bedrock_embedding_extra_body_cannot_replace_the_input_or_a_mapped_param
|
|||
assert json.loads(route.calls.last.request.read())[field] == expected
|
||||
|
||||
|
||||
def test_bedrock_embedding_signs_the_body_it_sends(respx_mock: respx.MockRouter) -> None:
|
||||
def test_bedrock_embedding_signs_the_body_it_sends(
|
||||
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
from botocore.auth import SigV4Auth
|
||||
from botocore.awsrequest import AWSRequest
|
||||
from botocore.credentials import Credentials
|
||||
|
||||
signing_time: Final = datetime(2026, 1, 2, 3, 4, 5, tzinfo=timezone.utc)
|
||||
monkeypatch.setattr("botocore.auth.get_current_datetime", lambda remove_tzinfo=True: signing_time)
|
||||
route = respx_mock.post(
|
||||
"https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-embed-text-v2:0/invoke"
|
||||
).mock(return_value=httpx.Response(200, json=titan_embedding_response))
|
||||
|
|
@ -2270,10 +2275,10 @@ def test_bedrock_embedding_signs_the_body_it_sends(respx_mock: respx.MockRouter)
|
|||
replay: Final = AWSRequest(
|
||||
method="POST", url=str(sent.url), data=sent.content, headers={"Content-Type": "application/json"}
|
||||
)
|
||||
replay.context["timestamp"] = sent.headers["x-amz-date"]
|
||||
SigV4Auth(
|
||||
Credentials("AKIAIOSFODNN7EXAMPLE", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), "bedrock", "us-east-1"
|
||||
).add_auth(replay)
|
||||
|
||||
assert sent.headers["x-amz-date"] == "20260102T030405Z"
|
||||
assert json.loads(sent.content)["dimensions"] == 256
|
||||
assert replay.headers["Authorization"] == sent.headers["authorization"]
|
||||
|
|
|
|||
|
|
@ -227,5 +227,5 @@ class TestInteractionByIdRequests:
|
|||
headers={},
|
||||
)
|
||||
|
||||
assert url == f"{GLOBAL_BASE}/abc123:cancel"
|
||||
assert url == f"{GLOBAL_BASE}/abc123/cancel"
|
||||
assert request_body == {}
|
||||
|
|
|
|||
|
|
@ -2003,6 +2003,38 @@ def test_get_model_from_request_nvidia_nim_relay_without_a_router_has_no_model()
|
|||
assert get_model_from_request(request_data=NIM_INFER_BODY, route="/nvidia_nim/nim-page-elements/v1/infer") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("route", "selector", "authorized_as"),
|
||||
[
|
||||
pytest.param("/v1beta/interactions/abc", "beta-id", "beta", id="deployment-id-authorizes-its-group"),
|
||||
pytest.param("/v1beta/interactions/abc/cancel", "beta", "beta", id="group-name-wins-over-a-same-named-id"),
|
||||
pytest.param("/interactions/abc", "no-such-id", "no-such-id", id="unknown-id-stays-as-sent"),
|
||||
pytest.param("/v1/files", "beta-id", "beta-id", id="other-routes-keep-the-raw-selector"),
|
||||
],
|
||||
)
|
||||
def test_get_model_from_request_authorizes_an_interaction_deployment_id_selector_as_its_group(
|
||||
route: str, selector: str, authorized_as: str
|
||||
) -> None:
|
||||
router: Final = Router(
|
||||
model_list=[
|
||||
{"model_name": "alpha", "litellm_params": {"model": "openai/gpt-6-astra"}, "model_info": {"id": "beta"}},
|
||||
{"model_name": "beta", "litellm_params": {"model": "openai/gpt-6-astra"}, "model_info": {"id": "beta-id"}},
|
||||
]
|
||||
)
|
||||
|
||||
for body, params, headers in (
|
||||
({}, {"model": selector}, None),
|
||||
({}, None, {"x-litellm-model": selector}),
|
||||
({"model": selector}, {"model": selector}, None),
|
||||
):
|
||||
assert (
|
||||
get_model_from_request(
|
||||
request_data=body, route=route, request_query_params=params, request_headers=headers, llm_router=router
|
||||
)
|
||||
== authorized_as
|
||||
)
|
||||
|
||||
|
||||
def test_get_model_from_request_includes_file_endpoint_header_model():
|
||||
assert (
|
||||
get_model_from_request(
|
||||
|
|
|
|||
|
|
@ -2,25 +2,37 @@
|
|||
Test for google_endpoints/endpoints.py
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from collections.abc import Iterator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import responses
|
||||
import respx
|
||||
from dotenv import load_dotenv
|
||||
from fastapi.testclient import TestClient
|
||||
from pydantic import JsonValue, ValidationError
|
||||
from starlette.requests import Request
|
||||
from starlette.responses import Response
|
||||
|
||||
import litellm
|
||||
from litellm.proxy._types import LiteLLM_TeamTableCachedObj, UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
|
||||
from litellm.proxy.google_endpoints.endpoints import _google_count_tokens_request, google_count_tokens
|
||||
from litellm.types.llms.vertex_ai import TokenCountDetailsResponse
|
||||
from tests._master_key import MASTER_KEY
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("system_key", "schema_key"),
|
||||
(("systemInstruction", "parameters"), ("system_instruction", "parametersJsonSchema")),
|
||||
)
|
||||
def test_google_count_tokens_translates_system_and_function_declarations(
|
||||
system_key: str, schema_key: str
|
||||
) -> None:
|
||||
def test_google_count_tokens_translates_system_and_function_declarations(system_key: str, schema_key: str) -> None:
|
||||
schema: Final = {"type": "object", "properties": {"query": {"type": "string"}}}
|
||||
request_body: Final[dict[str, JsonValue]] = {
|
||||
"contents": [{"role": "user", "parts": [{"text": "Look this up"}]}],
|
||||
|
|
@ -57,9 +69,7 @@ def test_google_count_tokens_translates_system_and_function_declarations(
|
|||
("parametersJsonSchema", {"type": "OBJECT", "properties": {"query": {"type": "STRING"}}}),
|
||||
),
|
||||
)
|
||||
def test_google_count_tokens_normalizes_parameters_only(
|
||||
schema_key: str, expected_schema: dict[str, JsonValue]
|
||||
) -> None:
|
||||
def test_google_count_tokens_normalizes_parameters_only(schema_key: str, expected_schema: dict[str, JsonValue]) -> None:
|
||||
request: Final = _google_count_tokens_request(
|
||||
"claude-sonnet",
|
||||
{
|
||||
|
|
@ -82,8 +92,8 @@ def test_google_count_tokens_normalizes_parameters_only(
|
|||
|
||||
assert request.tools == [{"name": "lookup", "input_schema": expected_schema}]
|
||||
|
||||
load_dotenv()
|
||||
|
||||
load_dotenv()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -195,3 +205,457 @@ def test_google_count_tokens_keeps_counting_system_text_that_sits_beside_media()
|
|||
assert _google_count_tokens_request(
|
||||
"counted", {**_CONTENTS, "systemInstruction": {"parts": [text_with_media]}}
|
||||
) == _google_count_tokens_request("counted", {**_CONTENTS, **_SYSTEM})
|
||||
|
||||
|
||||
_INTERACTION_ID: Final = "abc123"
|
||||
_VERTEX_GROUP: Final = "vertex-argon"
|
||||
_VERTEX_DEEP_RESEARCH_GROUP: Final = "vertex-dr"
|
||||
_DEEP_RESEARCH_AGENT: Final = "deep-research-preview-04-2026"
|
||||
_SPLIT_GROUP: Final = "vertex-split"
|
||||
_SHARED_GROUP: Final = "vertex-shared"
|
||||
_LITERAL_STAR_ACCESS_GROUP: Final = "vertex-sha*"
|
||||
_SPLIT_DEPLOYMENTS: Final = MappingProxyType({"split-east": "east-project", "split-west": "west-project"})
|
||||
_VERTEX_COLLECTION_URL: Final = (
|
||||
"https://aiplatform.googleapis.com/v1beta1/projects/deployment-project/locations/global/interactions"
|
||||
)
|
||||
_VERTEX_INTERACTION_URL: Final = f"{_VERTEX_COLLECTION_URL}/{_INTERACTION_ID}"
|
||||
_GEMINI_INTERACTION_URL: Final = f"https://generativelanguage.googleapis.com/v1beta/interactions/{_INTERACTION_ID}"
|
||||
_INTERACTION: Final = {"id": _INTERACTION_ID, "object": "interaction", "status": "cancelled"}
|
||||
_BY_ID_CALLS: Final = (
|
||||
("GET", f"/v1beta/interactions/{_INTERACTION_ID}", ""),
|
||||
("DELETE", f"/v1beta/interactions/{_INTERACTION_ID}", ""),
|
||||
("POST", f"/v1beta/interactions/{_INTERACTION_ID}/cancel", "/cancel"),
|
||||
)
|
||||
_BY_ID_ROUTES: Final = tuple((method, path) for method, path, _ in _BY_ID_CALLS)
|
||||
|
||||
|
||||
def _vertex_deployment(
|
||||
model_name: str,
|
||||
model: str,
|
||||
project: str = "deployment-project",
|
||||
deployment_id: str | None = None,
|
||||
order: int | None = None,
|
||||
access_groups: tuple[str, ...] = (),
|
||||
) -> Mapping[str, object]:
|
||||
credentials: Final = {
|
||||
"type": "authorized_user",
|
||||
"client_id": "deployment-client",
|
||||
"client_secret": "deployment-secret",
|
||||
"refresh_token": "deployment-refresh-token",
|
||||
"quota_project_id": "quota-project",
|
||||
}
|
||||
return {
|
||||
"model_name": model_name,
|
||||
"litellm_params": {
|
||||
"model": model,
|
||||
"vertex_project": project,
|
||||
"vertex_location": "global",
|
||||
"vertex_credentials": json.dumps(credentials),
|
||||
**({"order": order} if order is not None else {}),
|
||||
},
|
||||
**(
|
||||
{
|
||||
"model_info": {
|
||||
"id": deployment_id,
|
||||
**({"access_groups": list(access_groups)} if access_groups else {}),
|
||||
}
|
||||
}
|
||||
if deployment_id is not None
|
||||
else {}
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def google_oauth() -> Iterator[responses.RequestsMock]:
|
||||
with responses.RequestsMock(assert_all_requests_are_fired=False) as token_endpoint:
|
||||
token_endpoint.post(
|
||||
"https://oauth2.googleapis.com/token",
|
||||
json={"access_token": "deployment-token", "expires_in": 3600, "token_type": "Bearer"},
|
||||
)
|
||||
yield token_endpoint
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def interactions_proxy(
|
||||
httpx_transport: None, fake_provider_credentials: None, google_oauth: responses.RequestsMock
|
||||
) -> Iterator[TestClient]:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
router: Final = litellm.Router(
|
||||
model_list=[
|
||||
_vertex_deployment(_VERTEX_GROUP, "vertex_ai/gemini-4-argon"),
|
||||
_vertex_deployment(_VERTEX_DEEP_RESEARCH_GROUP, f"vertex_ai/{_DEEP_RESEARCH_AGENT}"),
|
||||
_vertex_deployment(_SPLIT_GROUP, "vertex_ai/gemini-4-argon", "east-project", "split-east", order=1),
|
||||
_vertex_deployment(_SPLIT_GROUP, "vertex_ai/gemini-4-argon", "west-project", "split-west", order=2),
|
||||
_vertex_deployment(
|
||||
_SHARED_GROUP,
|
||||
"vertex_ai/gemini-4-argon",
|
||||
"pub-project",
|
||||
"shared-pub",
|
||||
access_groups=("public", _LITERAL_STAR_ACCESS_GROUP),
|
||||
),
|
||||
_vertex_deployment(
|
||||
_SHARED_GROUP, "vertex_ai/gemini-4-argon", "priv-project", "shared-priv", access_groups=("private",)
|
||||
),
|
||||
]
|
||||
)
|
||||
proxy_server.cleanup_router_config_variables()
|
||||
try:
|
||||
with pytest.MonkeyPatch.context() as proxy_globals:
|
||||
proxy_globals.setattr(proxy_server, "llm_router", None)
|
||||
config_path: Final = Path(__file__).parents[1] / "test_configs" / "test_config_no_auth.yaml"
|
||||
asyncio.run(proxy_server.initialize(config=str(config_path)))
|
||||
config_router: Final = proxy_server.llm_router
|
||||
proxy_globals.setattr(proxy_server, "llm_router", router)
|
||||
client: Final = TestClient(proxy_server.app)
|
||||
try:
|
||||
yield client
|
||||
finally:
|
||||
client.close()
|
||||
if config_router is not None:
|
||||
config_router.retire()
|
||||
finally:
|
||||
router.retire()
|
||||
proxy_server.cleanup_router_config_variables()
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _cached_key(key: UserAPIKeyAuth, *teams: LiteLLM_TeamTableCachedObj) -> Iterator[str]:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
user_key: Final = "sk-cached-test-key"
|
||||
user_api_key_cache: Final = UserApiKeyCache()
|
||||
user_api_key_cache.set_cache(
|
||||
key=proxy_server.hash_token(user_key), value=key.model_copy(update={"token": proxy_server.hash_token(user_key)})
|
||||
)
|
||||
for team in teams:
|
||||
user_api_key_cache.set_cache(key=f"team_id:{team.team_id}", value=team)
|
||||
with pytest.MonkeyPatch.context() as key_auth:
|
||||
key_auth.setattr(proxy_server, "user_api_key_cache", user_api_key_cache)
|
||||
key_auth.setattr(proxy_server, "master_key", MASTER_KEY)
|
||||
key_auth.setattr(proxy_server, "prisma_client", "no-db-needed-for-a-cached-key")
|
||||
yield f"Bearer {user_key}"
|
||||
|
||||
|
||||
def _selected(selector: str) -> tuple[Mapping[str, str], Mapping[str, str]]:
|
||||
match selector:
|
||||
case "query":
|
||||
return {"model": _VERTEX_GROUP}, {}
|
||||
case "header":
|
||||
return {}, {"x-litellm-model": _VERTEX_GROUP}
|
||||
case _:
|
||||
raise AssertionError(selector)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("selector", ("query", "header"))
|
||||
@pytest.mark.parametrize(("method", "path", "upstream_suffix"), _BY_ID_CALLS)
|
||||
def test_interaction_by_id_reaches_the_named_vertex_deployment_with_its_project_and_credentials(
|
||||
interactions_proxy: TestClient,
|
||||
respx_mock: respx.MockRouter,
|
||||
selector: str,
|
||||
method: str,
|
||||
path: str,
|
||||
upstream_suffix: str,
|
||||
) -> None:
|
||||
upstream: Final = respx_mock.request(method, f"{_VERTEX_INTERACTION_URL}{upstream_suffix}").mock(
|
||||
return_value=httpx.Response(200, json=_INTERACTION)
|
||||
)
|
||||
|
||||
params, headers = _selected(selector)
|
||||
|
||||
response: Final = interactions_proxy.request(method, path, params=params, headers=headers)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json()["id"] == _INTERACTION_ID
|
||||
assert upstream.call_count == 1
|
||||
assert upstream.calls.last.request.headers["authorization"] == "Bearer deployment-token"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("query", ({}, {"model": ""}), ids=("no_selector", "empty_selector"))
|
||||
@pytest.mark.parametrize(("method", "path", "upstream_suffix"), _BY_ID_CALLS)
|
||||
def test_interaction_by_id_without_a_model_still_goes_to_google_ai_studio(
|
||||
interactions_proxy: TestClient,
|
||||
respx_mock: respx.MockRouter,
|
||||
query: Mapping[str, str],
|
||||
method: str,
|
||||
path: str,
|
||||
upstream_suffix: str,
|
||||
) -> None:
|
||||
upstream: Final = respx_mock.request(method, f"{_GEMINI_INTERACTION_URL}{upstream_suffix}").mock(
|
||||
return_value=httpx.Response(200, json=_INTERACTION)
|
||||
)
|
||||
|
||||
response: Final = interactions_proxy.request(method, path, params=query)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert upstream.call_count == 1
|
||||
assert upstream.calls.last.request.headers["x-goog-api-key"] == "unit-test"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("method", "path"), _BY_ID_ROUTES)
|
||||
def test_interaction_by_id_naming_an_unknown_model_group_is_refused_before_any_upstream_call(
|
||||
interactions_proxy: TestClient, respx_mock: respx.MockRouter, method: str, path: str
|
||||
) -> None:
|
||||
response: Final = interactions_proxy.request(method, path, params={"model": "no-such-group"})
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert "no-such-group" in response.text
|
||||
assert not respx_mock.calls
|
||||
|
||||
|
||||
@pytest.mark.parametrize("selector", ("query", "header"))
|
||||
@pytest.mark.parametrize(("method", "path"), _BY_ID_ROUTES)
|
||||
def test_interaction_by_id_naming_a_model_the_key_cannot_call_is_refused_before_any_upstream_call(
|
||||
interactions_proxy: TestClient,
|
||||
respx_mock: respx.MockRouter,
|
||||
selector: str,
|
||||
method: str,
|
||||
path: str,
|
||||
) -> None:
|
||||
params, headers = _selected(selector)
|
||||
|
||||
with _cached_key(UserAPIKeyAuth(models=["some-other-group"])) as authorization:
|
||||
response: Final = interactions_proxy.request(
|
||||
method, path, params=params, headers={**headers, "authorization": authorization}
|
||||
)
|
||||
|
||||
assert response.status_code == 403, response.text
|
||||
assert _VERTEX_GROUP in response.text
|
||||
assert not respx_mock.calls
|
||||
|
||||
|
||||
def _project_url(project: str) -> str:
|
||||
return f"https://aiplatform.googleapis.com/v1beta1/projects/{project}/locations/global/interactions"
|
||||
|
||||
|
||||
def _deployment_selector(selector: str, deployment_id: str) -> tuple[Mapping[str, str], Mapping[str, str]]:
|
||||
match selector:
|
||||
case "query":
|
||||
return {"model": deployment_id}, {}
|
||||
case "header":
|
||||
return {}, {"x-litellm-model": deployment_id}
|
||||
case _:
|
||||
raise AssertionError(selector)
|
||||
|
||||
|
||||
@pytest.mark.respx(assert_all_called=False)
|
||||
@pytest.mark.parametrize("selector", ("query", "header"))
|
||||
@pytest.mark.parametrize(("method", "path", "upstream_suffix"), _BY_ID_CALLS)
|
||||
def test_a_follow_up_by_the_created_deployment_id_reaches_the_project_that_holds_the_interaction(
|
||||
interactions_proxy: TestClient,
|
||||
respx_mock: respx.MockRouter,
|
||||
selector: str,
|
||||
method: str,
|
||||
path: str,
|
||||
upstream_suffix: str,
|
||||
) -> None:
|
||||
created: Final = MappingProxyType(
|
||||
{
|
||||
project: respx_mock.post(_project_url(project)).mock(
|
||||
return_value=httpx.Response(
|
||||
200, json={"id": _INTERACTION_ID, "object": "interaction", "status": "completed"}
|
||||
)
|
||||
)
|
||||
for project in _SPLIT_DEPLOYMENTS.values()
|
||||
}
|
||||
)
|
||||
followed_up: Final = MappingProxyType(
|
||||
{
|
||||
project: respx_mock.request(method, f"{_project_url(project)}/{_INTERACTION_ID}{upstream_suffix}").mock(
|
||||
return_value=httpx.Response(200, json=_INTERACTION)
|
||||
)
|
||||
for project in _SPLIT_DEPLOYMENTS.values()
|
||||
}
|
||||
)
|
||||
|
||||
create: Final = interactions_proxy.post("/v1beta/interactions", json={"model": _SPLIT_GROUP, "input": "Hi"})
|
||||
deployment_id: Final = create.headers["x-litellm-model-id"]
|
||||
project: Final = _SPLIT_DEPLOYMENTS[deployment_id]
|
||||
params, headers = _deployment_selector(selector, deployment_id)
|
||||
follow_up: Final = interactions_proxy.request(method, path, params=params, headers=headers)
|
||||
|
||||
assert create.status_code == 200, create.text
|
||||
assert created[project].call_count == 1
|
||||
assert follow_up.status_code == 200, follow_up.text
|
||||
assert {name: route.call_count for name, route in followed_up.items()} == {
|
||||
name: 1 if name == project else 0 for name in _SPLIT_DEPLOYMENTS.values()
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("selector", ("query", "header"))
|
||||
def test_a_key_allowed_the_group_reaches_the_named_deployment_although_the_group_prefers_another(
|
||||
interactions_proxy: TestClient, respx_mock: respx.MockRouter, selector: str
|
||||
) -> None:
|
||||
upstream: Final = respx_mock.get(f"{_project_url('west-project')}/{_INTERACTION_ID}").mock(
|
||||
return_value=httpx.Response(200, json=_INTERACTION)
|
||||
)
|
||||
params, headers = _deployment_selector(selector, "split-west")
|
||||
|
||||
with _cached_key(UserAPIKeyAuth(models=[_SPLIT_GROUP])) as authorization:
|
||||
response: Final = interactions_proxy.get(
|
||||
f"/v1beta/interactions/{_INTERACTION_ID}",
|
||||
params=params,
|
||||
headers={**headers, "authorization": authorization},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert upstream.call_count == 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("key", "teams", "deployment_id"),
|
||||
(
|
||||
pytest.param(UserAPIKeyAuth(models=[_VERTEX_DEEP_RESEARCH_GROUP]), (), "split-west", id="group-not-on-key"),
|
||||
pytest.param(
|
||||
UserAPIKeyAuth(team_id="team-without-split"),
|
||||
(LiteLLM_TeamTableCachedObj(team_id="team-without-split", models=[_VERTEX_DEEP_RESEARCH_GROUP]),),
|
||||
"split-west",
|
||||
id="group-not-on-team",
|
||||
),
|
||||
pytest.param(UserAPIKeyAuth(models=[_SPLIT_GROUP]), (), "no-such-deployment", id="unknown-deployment-id"),
|
||||
),
|
||||
)
|
||||
@pytest.mark.parametrize("selector", ("query", "header"))
|
||||
def test_a_follow_up_by_deployment_id_is_refused_when_the_key_cannot_call_its_group(
|
||||
interactions_proxy: TestClient,
|
||||
respx_mock: respx.MockRouter,
|
||||
google_oauth: responses.RequestsMock,
|
||||
key: UserAPIKeyAuth,
|
||||
teams: tuple[LiteLLM_TeamTableCachedObj, ...],
|
||||
deployment_id: str,
|
||||
selector: str,
|
||||
) -> None:
|
||||
params, headers = _deployment_selector(selector, deployment_id)
|
||||
|
||||
with _cached_key(key, *teams) as authorization:
|
||||
response: Final = interactions_proxy.get(
|
||||
f"/v1beta/interactions/{_INTERACTION_ID}",
|
||||
params=params,
|
||||
headers={**headers, "authorization": authorization},
|
||||
)
|
||||
|
||||
assert response.status_code == 403, response.text
|
||||
assert not respx_mock.calls
|
||||
assert not google_oauth.calls
|
||||
|
||||
|
||||
_PUBLIC_KEY: Final = UserAPIKeyAuth(models=["public"])
|
||||
_PUBLIC_TEAM_KEY: Final = UserAPIKeyAuth(team_id="team-public", team_models=["public"])
|
||||
_PUBLIC_TEAM: Final = LiteLLM_TeamTableCachedObj(team_id="team-public", models=["public"])
|
||||
_WILDCARD_TEAM_KEY: Final = UserAPIKeyAuth(team_id="team-wildcard", team_models=[f"{_SHARED_GROUP}*", "public"])
|
||||
_WILDCARD_TEAM: Final = LiteLLM_TeamTableCachedObj(team_id="team-wildcard", models=[f"{_SHARED_GROUP}*", "public"])
|
||||
|
||||
|
||||
@pytest.mark.parametrize("selector", ("query", "header"))
|
||||
@pytest.mark.parametrize(("method", "path"), _BY_ID_ROUTES)
|
||||
@pytest.mark.parametrize(
|
||||
("key", "teams"),
|
||||
(
|
||||
pytest.param(_PUBLIC_KEY, (), id="key-access-group"),
|
||||
pytest.param(_PUBLIC_TEAM_KEY, (_PUBLIC_TEAM,), id="team-access-group"),
|
||||
pytest.param(UserAPIKeyAuth(models=["vertex-sh*x", "public"]), (), id="near-match-wildcard"),
|
||||
pytest.param(UserAPIKeyAuth(models=["other-group*", "public"]), (), id="unrelated-wildcard"),
|
||||
pytest.param(UserAPIKeyAuth(models=[_LITERAL_STAR_ACCESS_GROUP]), (), id="access-group-named-like-a-wildcard"),
|
||||
),
|
||||
)
|
||||
def test_a_follow_up_pinned_to_a_deployment_outside_the_callers_access_group_is_refused(
|
||||
interactions_proxy: TestClient,
|
||||
respx_mock: respx.MockRouter,
|
||||
google_oauth: responses.RequestsMock,
|
||||
key: UserAPIKeyAuth,
|
||||
teams: tuple[LiteLLM_TeamTableCachedObj, ...],
|
||||
method: str,
|
||||
path: str,
|
||||
selector: str,
|
||||
) -> None:
|
||||
params, headers = _deployment_selector(selector, "shared-priv")
|
||||
|
||||
with _cached_key(key, *teams) as authorization:
|
||||
response: Final = interactions_proxy.request(
|
||||
method, path, params=params, headers={**headers, "authorization": authorization}
|
||||
)
|
||||
|
||||
assert response.status_code == 400, response.text
|
||||
assert not respx_mock.calls
|
||||
assert not google_oauth.calls
|
||||
|
||||
|
||||
@pytest.mark.parametrize("selector", ("query", "header"))
|
||||
@pytest.mark.parametrize(
|
||||
("key", "teams", "model", "project"),
|
||||
(
|
||||
pytest.param(_PUBLIC_KEY, (), "shared-pub", "pub-project", id="key-access-group-pins-its-deployment"),
|
||||
pytest.param(
|
||||
_PUBLIC_TEAM_KEY, (_PUBLIC_TEAM,), "shared-pub", "pub-project", id="team-access-group-pins-its-deployment"
|
||||
),
|
||||
pytest.param(_PUBLIC_KEY, (), _SHARED_GROUP, "pub-project", id="key-access-group-group-route"),
|
||||
pytest.param(
|
||||
_PUBLIC_TEAM_KEY, (_PUBLIC_TEAM,), _SHARED_GROUP, "pub-project", id="team-access-group-group-route"
|
||||
),
|
||||
pytest.param(
|
||||
UserAPIKeyAuth(models=[_SHARED_GROUP]), (), "shared-priv", "priv-project", id="owning-group-pins-any"
|
||||
),
|
||||
pytest.param(
|
||||
UserAPIKeyAuth(models=["private"]), (), "shared-priv", "priv-project", id="private-access-group-pins-its"
|
||||
),
|
||||
pytest.param(
|
||||
UserAPIKeyAuth(models=[f"{_SHARED_GROUP}*", "public"]),
|
||||
(),
|
||||
"shared-priv",
|
||||
"priv-project",
|
||||
id="key-wildcard-beside-access-group-pins-any",
|
||||
),
|
||||
pytest.param(
|
||||
_WILDCARD_TEAM_KEY,
|
||||
(_WILDCARD_TEAM,),
|
||||
"shared-priv",
|
||||
"priv-project",
|
||||
id="team-wildcard-beside-access-group-pins-any",
|
||||
),
|
||||
pytest.param(
|
||||
UserAPIKeyAuth(models=[f"{_SHARED_GROUP}*"]),
|
||||
(),
|
||||
"shared-priv",
|
||||
"priv-project",
|
||||
id="wildcard-alone-pins-any",
|
||||
),
|
||||
pytest.param(
|
||||
UserAPIKeyAuth(models=[f"{_SHARED_GROUP}*", "public"]),
|
||||
(),
|
||||
_SHARED_GROUP,
|
||||
"pub-project",
|
||||
id="wildcard-beside-access-group-group-route-unchanged",
|
||||
),
|
||||
pytest.param(
|
||||
UserAPIKeyAuth(models=[_LITERAL_STAR_ACCESS_GROUP]),
|
||||
(),
|
||||
"shared-pub",
|
||||
"pub-project",
|
||||
id="access-group-named-like-a-wildcard-pins-its-deployment",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_a_follow_up_reaches_a_deployment_the_callers_grant_covers(
|
||||
interactions_proxy: TestClient,
|
||||
respx_mock: respx.MockRouter,
|
||||
key: UserAPIKeyAuth,
|
||||
teams: tuple[LiteLLM_TeamTableCachedObj, ...],
|
||||
model: str,
|
||||
project: str,
|
||||
selector: str,
|
||||
) -> None:
|
||||
upstream: Final = respx_mock.get(f"{_project_url(project)}/{_INTERACTION_ID}").mock(
|
||||
return_value=httpx.Response(200, json=_INTERACTION)
|
||||
)
|
||||
params, headers = _deployment_selector(selector, model)
|
||||
|
||||
with _cached_key(key, *teams) as authorization:
|
||||
response: Final = interactions_proxy.get(
|
||||
f"/v1beta/interactions/{_INTERACTION_ID}",
|
||||
params=params,
|
||||
headers={**headers, "authorization": authorization},
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert upstream.call_count == 1
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue