diff --git a/litellm/llms/vertex_ai/interactions/transformation.py b/litellm/llms/vertex_ai/interactions/transformation.py index b8e13059917..d81818d021d 100644 --- a/litellm/llms/vertex_ai/interactions/transformation.py +++ b/litellm/llms/vertex_ai/interactions/transformation.py @@ -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") diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 7d0ca935cde..72dbe297c98 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -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, diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 31b88e8110c..df0c91dfde2 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -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, diff --git a/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py b/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py index 63526e917a6..b128ad1a564 100644 --- a/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py @@ -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"] diff --git a/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py b/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py index 3364b1b3872..bebc7d8e021 100644 --- a/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py +++ b/tests/unit/llms/vertex_ai/interactions/test_vertex_ai_interactions_transformation.py @@ -227,5 +227,5 @@ class TestInteractionByIdRequests: headers={}, ) - assert url == f"{GLOBAL_BASE}/abc123:cancel" + assert url == f"{GLOBAL_BASE}/abc123/cancel" assert request_body == {} diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index 6a7bbb2f0ce..bcdacd0996e 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -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( diff --git a/tests/unit/proxy/google_endpoints/test_endpoints.py b/tests/unit/proxy/google_endpoints/test_endpoints.py index 8e1a39dd176..e6df7bf1d4f 100644 --- a/tests/unit/proxy/google_endpoints/test_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_endpoints.py @@ -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