import uuid from collections.abc import Iterator from typing import Final import httpx import pytest from integration._support.client import Gateway, Scenario, object_value, string_value from integration._support.database import read_rows from pydantic import BaseModel, JsonValue class _ObservedRequest(BaseModel): body: dict[str, JsonValue] class _Observations(BaseModel): requests: list[_ObservedRequest] class _TeamState(BaseModel): team_id: str models: list[str] | None = None class _TeamModelRow(BaseModel): models: list[str] | None = None class _TeamInfo(BaseModel): team_info: _TeamState class _ProxyError(BaseModel): type: str class _ErrorResponse(BaseModel): error: _ProxyError @pytest.fixture def upstream(gateway: Gateway) -> Iterator[httpx.Client]: with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as client: client.get("/__observations").raise_for_status() yield client def _observed_requests(upstream: httpx.Client) -> list[JsonValue]: observed: Final = upstream.get("/__observations") observed.raise_for_status() requests: Final = object_value(observed.json())["requests"] assert isinstance(requests, list) return requests def _calls_to(observed: list[JsonValue], provider_model: str) -> int: return sum(object_value(object_value(request)["body"]).get("model") == provider_model for request in observed) def _observed_provider_models(upstream: httpx.Client) -> tuple[str, ...]: observed: Final = upstream.get("/__observations") assert observed.status_code == 200, observed.text response: Final = _Observations.model_validate_json(observed.text) return tuple(string_value(request.body["model"]) for request in response.requests) def _chat(gateway: Gateway, model: str, key: str) -> httpx.Response: return gateway.request( "POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": f"team model {uuid.uuid4().hex}"}]}, key=key, ) def _team_model(scenario: Scenario) -> tuple[str, str]: provider_model: Final = f"team-model-{uuid.uuid4().hex}" return scenario.model(model=f"openai/{provider_model}"), provider_model def _team_info(gateway: Gateway, team_id: str) -> _TeamInfo: response: Final = gateway.request("GET", "/team/info", params={"team_id": team_id}) assert response.status_code == 200, response.text return _TeamInfo.model_validate_json(response.text) def _stored_team_models(team_id: str) -> _TeamModelRow: rows: Final = read_rows('SELECT models FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,)) assert len(rows) == 1, repr(rows) return _TeamModelRow.model_validate(rows[0]) def _model_denied(response: httpx.Response) -> None: assert response.status_code == 403, response.text assert _ErrorResponse.model_validate_json(response.text).error.type == "team_model_access_denied", response.text def _ids(listing: dict[str, JsonValue]) -> set[JsonValue]: data: Final = listing["data"] assert isinstance(data, list) return {object_value(entry)["id"] for entry in data} def test_a_model_created_for_a_team_is_listed_in_that_teams_models(gateway: Gateway) -> None: with gateway.scenario() as scenario: team: Final = scenario.team(models=[]) other: Final = scenario.team(models=[]) model: Final = scenario.model(model_info={"team_id": team}) own_models: Final = object_value(gateway.get("/team/info", {"team_id": team})["team_info"])["models"] other_models: Final = object_value(gateway.get("/team/info", {"team_id": other})["team_info"])["models"] assert isinstance(own_models, list) and isinstance(other_models, list) assert model in own_models assert model not in other_models def test_a_team_model_is_listed_and_served_only_for_keys_of_its_team(gateway: Gateway, upstream: httpx.Client) -> None: with gateway.scenario() as scenario: team: Final = scenario.team(models=[]) other: Final = scenario.team(models=[]) provider_model: Final = f"team-scoped-{uuid.uuid4().hex}" model: Final = scenario.model(model=f"openai/{provider_model}", model_info={"team_id": team}) team_key: Final = scenario.key(team_id=team) other_key: Final = scenario.key(team_id=other) assert model in _ids(object_value(gateway.request("GET", "/models", key=team_key).json())) assert model not in _ids(object_value(gateway.request("GET", "/models", key=other_key).json())) served: Final = _chat(gateway, model, team_key) assert served.status_code == 200, served.text refused: Final = _chat(gateway, model, other_key) assert refused.status_code == 400, refused.text observed: Final = _observed_requests(upstream) assert _calls_to(observed, provider_model) == 1, observed def _v2_team_public_names(gateway: Gateway, key: str, model: str) -> list[JsonValue]: response: Final = gateway.request("GET", "/v2/model/info", key=key, params={"model": model}) assert response.status_code == 200, response.text return [entry["model_info"].get("team_public_model_name") for entry in response.json()["data"]] def test_v2_model_info_reports_a_team_model_to_team_and_non_team_keys(gateway: Gateway) -> None: with gateway.scenario() as scenario: team: Final = scenario.team(models=[]) model: Final = scenario.model(model_info={"team_id": team}) assert _v2_team_public_names(gateway, scenario.key(team_id=team), model) == [model] assert _v2_team_public_names(gateway, scenario.key(), model) == [model] @pytest.mark.parametrize("set_on", ["team_new", "team_update"]) def test_team_model_alias_routes_a_team_key_to_its_target( gateway: Gateway, upstream: httpx.Client, set_on: str ) -> None: with gateway.scenario() as scenario: provider_model: Final = f"team-alias-{uuid.uuid4().hex}" model: Final = scenario.model(model=f"openai/{provider_model}") alias: Final = f"alias-{uuid.uuid4().hex}" team: Final = ( scenario.team(models=[model], model_aliases={alias: model}) if set_on == "team_new" else scenario.team(models=[model]) ) if set_on == "team_update": gateway.post("/team/update", {"team_id": team, "model_aliases": {alias: model}}) key: Final = scenario.key(team_id=team, models=[model]) response: Final = _chat(gateway, alias, key) assert response.status_code == 200, response.text assert string_value(response.json()["model"]) == alias observed: Final = _observed_requests(upstream) assert _calls_to(observed, provider_model) == 1, observed unaliased: Final = _chat(gateway, f"alias-{uuid.uuid4().hex}", key) assert unaliased.status_code == 403, unaliased.text assert unaliased.json()["error"]["type"] == "key_model_access_denied" after_refusal: Final = _observed_requests(upstream) assert _calls_to(after_refusal, provider_model) == 0, after_refusal def test_team_model_add_and_delete_change_what_a_warmed_team_key_can_call( gateway: Gateway, upstream: httpx.Client ) -> None: with gateway.scenario() as scenario: model_one, provider_one = _team_model(scenario) model_two, provider_two = _team_model(scenario) model_three, provider_three = _team_model(scenario) team: Final = scenario.team(models=[model_one]) key: Final = scenario.key(team_id=team) first: Final = _chat(gateway, model_one, key) assert first.status_code == 200, first.text initially_denied: Final = _chat(gateway, model_two, key) _model_denied(initially_denied) observed_before_add: Final = _observed_provider_models(upstream) assert observed_before_add == (provider_one,), repr(observed_before_add) added: Final = gateway.request("POST", "/team/model/add", {"team_id": team, "models": [model_two, model_three]}) assert added.status_code == 200, added.text added_state: Final = _TeamState.model_validate_json(added.text) expected_added: Final = sorted((model_one, model_two, model_three)) assert sorted(added_state.models or ()) == expected_added, added.text added_info: Final = _team_info(gateway, team) assert sorted(added_info.team_info.models or ()) == expected_added, repr(added_info) added_row: Final = _stored_team_models(team) assert sorted(added_row.models or ()) == expected_added, repr(added_row) second: Final = _chat(gateway, model_two, key) assert second.status_code == 200, second.text third: Final = _chat(gateway, model_three, key) assert third.status_code == 200, third.text deleted: Final = gateway.request("POST", "/team/model/delete", {"team_id": team, "models": [model_two]}) assert deleted.status_code == 200, deleted.text deleted_state: Final = _TeamState.model_validate_json(deleted.text) expected_remaining: Final = sorted((model_one, model_three)) assert sorted(deleted_state.models or ()) == expected_remaining, deleted.text deleted_info: Final = _team_info(gateway, team) assert sorted(deleted_info.team_info.models or ()) == expected_remaining, repr(deleted_info) deleted_row: Final = _stored_team_models(team) assert sorted(deleted_row.models or ()) == expected_remaining, repr(deleted_row) denied_after_delete: Final = _chat(gateway, model_two, key) _model_denied(denied_after_delete) still_allowed_one: Final = _chat(gateway, model_one, key) assert still_allowed_one.status_code == 200, still_allowed_one.text still_allowed_three: Final = _chat(gateway, model_three, key) assert still_allowed_three.status_code == 200, still_allowed_three.text observed_after_add_and_delete: Final = _observed_provider_models(upstream) assert observed_before_add + observed_after_add_and_delete == ( provider_one, provider_two, provider_three, provider_one, provider_three, ), repr((observed_before_add, observed_after_add_and_delete)) def test_team_model_add_on_an_unrestricted_team_keeps_every_other_model( gateway: Gateway, upstream: httpx.Client ) -> None: with gateway.scenario() as scenario: model_one, provider_one = _team_model(scenario) model_two, provider_two = _team_model(scenario) team: Final = scenario.team(models=[]) key: Final = scenario.key(team_id=team) before_add: Final = _chat(gateway, model_two, key) assert before_add.status_code == 200, before_add.text added: Final = gateway.request("POST", "/team/model/add", {"team_id": team, "models": [model_one]}) assert added.status_code == 200, added.text expected: Final = sorted(("all-proxy-models", model_one)) added_state: Final = _TeamState.model_validate_json(added.text) assert sorted(added_state.models or ()) == expected, added.text added_info: Final = _team_info(gateway, team) assert sorted(added_info.team_info.models or ()) == expected, repr(added_info) added_row: Final = _stored_team_models(team) assert sorted(added_row.models or ()) == expected, repr(added_row) still_unrestricted: Final = _chat(gateway, model_two, key) assert still_unrestricted.status_code == 200, still_unrestricted.text added_model: Final = _chat(gateway, model_one, key) assert added_model.status_code == 200, added_model.text observed_after_add: Final = _observed_provider_models(upstream) assert observed_after_add == (provider_two, provider_two, provider_one), repr(observed_after_add)