mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
test(integration): cover model_management endpoint contracts (#44684)
* test(integration): cover model_management endpoint contracts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): pin the exact model access refusal after a rename Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): wait for allowlist convergence in the rename test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): give the rename test lint suppression its reason Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): close reviewer gaps on model management contracts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): share the blocked serving assertions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): strengthen authorization and configuration regression tests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): converge spend readbacks, ignore model discovery, cover slash credential names Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): pin the dashboard credential edit body round-trip Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mateo <mateo@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
925c0a9586
commit
013c440597
5 changed files with 1514 additions and 28 deletions
429
tests/integration/authorization/test_access_group_budget.py
Normal file
429
tests/integration/authorization/test_access_group_budget.py
Normal file
|
|
@ -0,0 +1,429 @@
|
|||
"""A model access group's shared budget: set and read back, enforced on real chat traffic, and cleared.
|
||||
|
||||
The scripted provider bills every call at 10 prompt and 10 completion tokens, so at the rates below one
|
||||
call costs 0.00011 and exhausts the 0.0001 budget the group is given.
|
||||
"""
|
||||
|
||||
import json
|
||||
import threading
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Final, Literal
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.process import owned_proxy
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
|
||||
_INPUT_RATE: Final = 0.000001
|
||||
_OUTPUT_RATE: Final = 0.00001
|
||||
_MAX_BUDGET: Final = 0.0001
|
||||
_CALL_COST: Final = 10 * _INPUT_RATE + 10 * _OUTPUT_RATE
|
||||
_PROVIDER_MODEL: Final = "group-budget-probe"
|
||||
Grant = Literal["key", "team", "organization", "project"]
|
||||
|
||||
|
||||
class _Budget(BaseModel):
|
||||
budget_id: str
|
||||
max_budget: float | None
|
||||
soft_budget: float | None
|
||||
budget_duration: str | None
|
||||
|
||||
|
||||
class _GroupBudget(BaseModel):
|
||||
access_group: str
|
||||
spend: float
|
||||
budget: _Budget | None
|
||||
|
||||
|
||||
class _DeletedBudget(BaseModel):
|
||||
access_group: str
|
||||
budget_deleted: bool
|
||||
message: str
|
||||
|
||||
|
||||
class _Usage(BaseModel):
|
||||
prompt_tokens: int
|
||||
completion_tokens: int
|
||||
total_tokens: int
|
||||
|
||||
|
||||
class _Completion(BaseModel):
|
||||
id: str
|
||||
usage: _Usage
|
||||
|
||||
|
||||
class _StreamDelta(BaseModel):
|
||||
content: str | None = None
|
||||
|
||||
|
||||
class _StreamChoice(BaseModel):
|
||||
delta: _StreamDelta
|
||||
|
||||
|
||||
class _StreamChunk(BaseModel):
|
||||
choices: list[_StreamChoice]
|
||||
usage: _Usage | None = None
|
||||
|
||||
|
||||
class _ErrorBody(BaseModel):
|
||||
message: str
|
||||
type: str
|
||||
param: str | None
|
||||
code: str
|
||||
|
||||
|
||||
class _Error(BaseModel):
|
||||
error: _ErrorBody
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Group:
|
||||
name: str
|
||||
model: str
|
||||
key: str
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _BudgetProxy:
|
||||
gateway: Gateway
|
||||
reserved: bool
|
||||
|
||||
|
||||
def _is_model_discovery(request: Request) -> bool:
|
||||
return (request.method, request.target) == ("GET", "/v1/models")
|
||||
|
||||
|
||||
def _provider_calls(wire: Wire) -> tuple[Request, ...]:
|
||||
return tuple(entry for entry in wire.drain() if not _is_model_discovery(entry))
|
||||
|
||||
|
||||
def _billed(request: Request) -> Reply:
|
||||
if _is_model_discovery(request):
|
||||
return Reply(body=b'{"object":"list","data":[]}')
|
||||
assert (request.method, request.target) == ("POST", "/v1/chat/completions"), request
|
||||
body: Final = json.loads(request.body)
|
||||
if body.get("stream") is True:
|
||||
identity: Final = f"chatcmpl-{uuid.uuid4().hex}"
|
||||
frames: Final = (
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 0,
|
||||
"model": _PROVIDER_MODEL,
|
||||
"choices": [{"index": 0, "delta": {"role": "assistant", "content": "billed"}, "finish_reason": None}],
|
||||
},
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 0,
|
||||
"model": _PROVIDER_MODEL,
|
||||
"choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}],
|
||||
},
|
||||
{
|
||||
"id": identity,
|
||||
"object": "chat.completion.chunk",
|
||||
"created": 0,
|
||||
"model": _PROVIDER_MODEL,
|
||||
"choices": [],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 10, "total_tokens": 20},
|
||||
},
|
||||
)
|
||||
chunks: Final = tuple(f"data: {json.dumps(frame)}\n\n".encode() for frame in frames) + (b"data: [DONE]\n\n",)
|
||||
return Reply(content_type="text/event-stream", chunks=chunks)
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": f"chatcmpl-{uuid.uuid4().hex}",
|
||||
"object": "chat.completion",
|
||||
"created": 0,
|
||||
"model": _PROVIDER_MODEL,
|
||||
"choices": [
|
||||
{"index": 0, "message": {"role": "assistant", "content": "billed"}, "finish_reason": "stop"}
|
||||
],
|
||||
"usage": {"prompt_tokens": 10, "completion_tokens": 10, "total_tokens": 20},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture(
|
||||
params=("reserved", "unreserved"),
|
||||
ids=("reservation", "no-reservation"),
|
||||
scope="module",
|
||||
)
|
||||
def proxy(request: pytest.FixtureRequest, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_BudgetProxy]:
|
||||
"""The shared proxy reserves budget per request; the owned one runs with ``disable_budget_reservation``."""
|
||||
with gateway_from_environment() as gateway:
|
||||
if request.param == "reserved":
|
||||
yield _BudgetProxy(gateway=gateway, reserved=True)
|
||||
return
|
||||
directory: Final = tmp_path_factory.mktemp("access-group-budget-unreserved")
|
||||
loaded_config: Final = TypeAdapter(dict[str, object]).validate_python(
|
||||
yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())
|
||||
)
|
||||
general_settings: Final = TypeAdapter(dict[str, object]).validate_python(
|
||||
loaded_config.get("general_settings", {})
|
||||
)
|
||||
config: Final = {
|
||||
**loaded_config,
|
||||
"general_settings": {**general_settings, "disable_budget_reservation": True},
|
||||
}
|
||||
path: Final = directory / "unreserved.yaml"
|
||||
path.write_text(yaml.safe_dump(config))
|
||||
with owned_proxy(gateway, directory, {}, config=path) as owned:
|
||||
yield _BudgetProxy(gateway=owned, reserved=False)
|
||||
|
||||
|
||||
def _group_key(scenario: Scenario, name: str, grant: Grant) -> str:
|
||||
match grant:
|
||||
case "key":
|
||||
return scenario.key(models=[name])
|
||||
case "team":
|
||||
team_id: Final = scenario.team(models=[name])
|
||||
return scenario.key(team_id=team_id, models=[])
|
||||
case "organization":
|
||||
organization: Final = scenario.organization(models=[name])
|
||||
return scenario.key(organization_id=organization, models=[])
|
||||
case "project":
|
||||
project_team: Final = scenario.team()
|
||||
project_id: Final = scenario.project(project_team, models=[name])
|
||||
return scenario.key(team_id=project_team, project_id=project_id, models=[])
|
||||
|
||||
|
||||
def _budgeted_group(scenario: Scenario, wire: Wire, grant: Grant = "key") -> _Group:
|
||||
"""A database model in a fresh access group with one grant and a 0.0001 group budget."""
|
||||
name: Final = f"integration-group-{uuid.uuid4().hex}"
|
||||
model: Final = scenario.model(
|
||||
model=f"openai/{_PROVIDER_MODEL}",
|
||||
api_base=f"{wire.url}/v1",
|
||||
input_cost_per_token=_INPUT_RATE,
|
||||
output_cost_per_token=_OUTPUT_RATE,
|
||||
model_info={"access_groups": [name]},
|
||||
)
|
||||
key: Final = _group_key(scenario, name, grant)
|
||||
gateway: Final = scenario.gateway
|
||||
created: Final = gateway.request(
|
||||
"PUT", f"/access_group/{name}/budget", {"max_budget": _MAX_BUDGET, "budget_duration": "30d"}
|
||||
)
|
||||
assert created.status_code == 200, created.text
|
||||
stored: Final = _GroupBudget.model_validate_json(created.content)
|
||||
assert stored.budget is not None, created.text
|
||||
scenario.cleanups.callback(scenario.delete_budget, stored.budget.budget_id)
|
||||
scenario.cleanups.callback(gateway.request, "DELETE", f"/access_group/{name}/budget")
|
||||
assert stored == _GroupBudget(
|
||||
access_group=name,
|
||||
spend=0.0,
|
||||
budget=_Budget(
|
||||
budget_id=stored.budget.budget_id, max_budget=_MAX_BUDGET, soft_budget=None, budget_duration="30d"
|
||||
),
|
||||
), created.text
|
||||
read_back: Final = gateway.request("GET", f"/access_group/{name}/budget")
|
||||
assert read_back.status_code == 200, read_back.text
|
||||
assert _GroupBudget.model_validate_json(read_back.content) == stored, read_back.text
|
||||
assert _budget_rows(name) == [
|
||||
{"budget_id": stored.budget.budget_id, "spend": 0.0, "max_budget": _MAX_BUDGET, "budget_duration": "30d"}
|
||||
]
|
||||
return _Group(name, model, key)
|
||||
|
||||
|
||||
def _budget_rows(group: str) -> list[dict[str, object]]:
|
||||
return read_rows(
|
||||
'SELECT g.budget_id, g.spend, b.max_budget, b.budget_duration FROM "LiteLLM_ModelAccessGroupBudgetTable" g '
|
||||
'LEFT JOIN "LiteLLM_BudgetTable" b ON b.budget_id = g.budget_id WHERE g.access_group_name = %s',
|
||||
(group,),
|
||||
)
|
||||
|
||||
|
||||
def _chat(gateway: Gateway, group: _Group, text: str, max_tokens: int, spelling: str = "json") -> httpx.Response:
|
||||
body: Final = {
|
||||
"model": group.model,
|
||||
"messages": [{"role": "user", "content": text}],
|
||||
"max_tokens": max_tokens,
|
||||
**({"stream": True, "stream_options": {"include_usage": True}} if spelling == "stream" else {}),
|
||||
}
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
body,
|
||||
key=group.key,
|
||||
)
|
||||
|
||||
|
||||
def _sent(text: str, max_tokens: int, spelling: str = "json") -> dict[str, object]:
|
||||
return {
|
||||
"model": _PROVIDER_MODEL,
|
||||
"messages": [{"role": "user", "content": text}],
|
||||
"max_tokens": max_tokens,
|
||||
**({"stream": True, "stream_options": {"include_usage": True}} if spelling == "stream" else {}),
|
||||
}
|
||||
|
||||
|
||||
def _assert_served(response: httpx.Response, wire: Wire, text: str, max_tokens: int, spelling: str = "json") -> None:
|
||||
assert response.status_code == 200, response.text
|
||||
if spelling == "stream":
|
||||
lines: Final = tuple(line for line in response.text.splitlines() if line)
|
||||
frames: Final = tuple(
|
||||
_StreamChunk.model_validate_json(line.removeprefix("data: "))
|
||||
for line in lines
|
||||
if line.startswith("data: ") and line != "data: [DONE]"
|
||||
)
|
||||
assert "".join(choice.delta.content or "" for frame in frames for choice in frame.choices) == "billed", (
|
||||
response.text
|
||||
)
|
||||
assert frames[-1].usage == _Usage(prompt_tokens=10, completion_tokens=10, total_tokens=20), response.text
|
||||
assert lines[-1] == "data: [DONE]", response.text
|
||||
else:
|
||||
assert _Completion.model_validate_json(response.content).usage == _Usage(
|
||||
prompt_tokens=10, completion_tokens=10, total_tokens=20
|
||||
), response.text
|
||||
received: Final = _provider_calls(wire)
|
||||
assert [json.loads(entry.body) for entry in received] == [_sent(text, max_tokens, spelling)], response.text
|
||||
assert [entry.headers["authorization"] for entry in received] == ["Bearer integration-provider-key"]
|
||||
|
||||
|
||||
def _assert_refused_for_budget(response: httpx.Response, wire: Wire, message: str) -> None:
|
||||
assert response.status_code == 422, response.text
|
||||
assert _Error.model_validate_json(response.content).error == _ErrorBody(
|
||||
message=message, type="budget_exceeded", param=None, code="422"
|
||||
), response.text
|
||||
assert _provider_calls(wire) == (), f"A refused request reached the provider: {response.text}"
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
@pytest.mark.parametrize("spelling", ("json", "stream"))
|
||||
def test_access_group_budget_refuses_traffic_once_spent_and_serves_again_after_delete(
|
||||
proxy: _BudgetProxy, spelling: str
|
||||
) -> None:
|
||||
gateway: Final = proxy.gateway
|
||||
arrived: Final = threading.Event()
|
||||
release: Final = threading.Event()
|
||||
|
||||
def held(request: Request) -> Reply:
|
||||
if proxy.reserved and spelling == "stream" and b"first call under budget" in request.body:
|
||||
arrived.set()
|
||||
assert release.wait(timeout=30), "The held call was never released"
|
||||
return _billed(request)
|
||||
|
||||
with wire_server(held) as wire, gateway.scenario() as scenario:
|
||||
group: Final = _budgeted_group(scenario, wire)
|
||||
if proxy.reserved and spelling == "stream":
|
||||
with ThreadPoolExecutor(max_workers=1) as pool:
|
||||
in_flight: Final = pool.submit(_chat, gateway, group, "first call under budget", 1000, spelling)
|
||||
try:
|
||||
assert arrived.wait(timeout=30), "The held call never reached the provider"
|
||||
concurrent: Final = _chat(gateway, group, "concurrent call", 1, spelling)
|
||||
finally:
|
||||
release.set()
|
||||
_assert_served(in_flight.result(timeout=30), wire, "first call under budget", 1000, spelling)
|
||||
_assert_refused_for_budget(
|
||||
concurrent,
|
||||
wire,
|
||||
f"Budget has been exceeded! Model access group={group.name} Current cost: {_MAX_BUDGET}, "
|
||||
f"Max budget: {_MAX_BUDGET}",
|
||||
)
|
||||
else:
|
||||
_assert_served(
|
||||
_chat(gateway, group, "first call under budget", 1, spelling),
|
||||
wire,
|
||||
"first call under budget",
|
||||
1,
|
||||
spelling,
|
||||
)
|
||||
eventually(
|
||||
lambda: _GroupBudget.model_validate_json(
|
||||
gateway.request("GET", f"/access_group/{group.name}/budget").content
|
||||
),
|
||||
lambda value: value.spend == pytest.approx(_CALL_COST),
|
||||
seconds=60,
|
||||
)
|
||||
assert _budget_rows(group.name)[0]["spend"] == pytest.approx(_CALL_COST)
|
||||
|
||||
_assert_refused_for_budget(
|
||||
_chat(gateway, group, "call over budget", 1, spelling),
|
||||
wire,
|
||||
f"Budget has been exceeded! Model access group={group.name} Current cost: {_CALL_COST}, Max budget: {_MAX_BUDGET}",
|
||||
)
|
||||
|
||||
deleted: Final = gateway.request("DELETE", f"/access_group/{group.name}/budget")
|
||||
assert deleted.status_code == 200, deleted.text
|
||||
assert _DeletedBudget.model_validate_json(deleted.content) == _DeletedBudget(
|
||||
access_group=group.name,
|
||||
budget_deleted=True,
|
||||
message=f"Budget for access group '{group.name}' deleted successfully",
|
||||
), deleted.text
|
||||
cleared: Final = gateway.request("GET", f"/access_group/{group.name}/budget")
|
||||
assert _GroupBudget.model_validate_json(cleared.content) == _GroupBudget(
|
||||
access_group=group.name, spend=0.0, budget=None
|
||||
), cleared.text
|
||||
assert _budget_rows(group.name) == []
|
||||
_assert_served(
|
||||
_chat(gateway, group, "call after delete", 1, spelling),
|
||||
wire,
|
||||
"call after delete",
|
||||
1,
|
||||
spelling,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.timeout(120)
|
||||
@pytest.mark.parametrize("grant", ("team", "organization", "project"))
|
||||
def test_access_group_budget_refuses_a_key_granted_the_group_through_its_scope(
|
||||
proxy: _BudgetProxy, grant: Literal["team", "organization", "project"]
|
||||
) -> None:
|
||||
gateway: Final = proxy.gateway
|
||||
text: Final = f"group granted through {grant}"
|
||||
with wire_server(_billed) as wire, gateway.scenario() as scenario:
|
||||
group: Final = _budgeted_group(scenario, wire, grant)
|
||||
_assert_served(_chat(gateway, group, text, 1), wire, text, 1)
|
||||
eventually(
|
||||
lambda: _GroupBudget.model_validate_json(
|
||||
gateway.request("GET", f"/access_group/{group.name}/budget").content
|
||||
),
|
||||
lambda value: value.spend == pytest.approx(_CALL_COST),
|
||||
seconds=60,
|
||||
)
|
||||
assert _provider_calls(wire) == (), "The served request was not drained before the refusal"
|
||||
_assert_refused_for_budget(
|
||||
_chat(gateway, group, f"{text} over budget", 1),
|
||||
wire,
|
||||
f"Budget has been exceeded! Model access group={group.name} Current cost: {_CALL_COST}, "
|
||||
f"Max budget: {_MAX_BUDGET}",
|
||||
)
|
||||
assert _budget_rows(group.name)[0]["spend"] == pytest.approx(_CALL_COST)
|
||||
|
||||
|
||||
@pytest.mark.timeout(120)
|
||||
def test_reservation_refuses_a_concurrent_call_while_an_in_flight_call_holds_the_group_budget(
|
||||
gateway: Gateway,
|
||||
) -> None:
|
||||
arrived: Final = threading.Event()
|
||||
release: Final = threading.Event()
|
||||
|
||||
def held(request: Request) -> Reply:
|
||||
if b"held call" in request.body:
|
||||
arrived.set()
|
||||
assert release.wait(timeout=30), "The held call was never released"
|
||||
return _billed(request)
|
||||
|
||||
with wire_server(held) as wire, gateway.scenario() as scenario, ThreadPoolExecutor(max_workers=1) as pool:
|
||||
group: Final = _budgeted_group(scenario, wire)
|
||||
in_flight: Final = pool.submit(_chat, gateway, group, "held call", 1000)
|
||||
try:
|
||||
assert arrived.wait(timeout=30), "The held call never reached the provider"
|
||||
concurrent: Final = _chat(gateway, group, "concurrent call", 1)
|
||||
finally:
|
||||
release.set()
|
||||
_assert_served(in_flight.result(timeout=30), wire, "held call", 1000)
|
||||
_assert_refused_for_budget(
|
||||
concurrent,
|
||||
wire,
|
||||
f"Budget has been exceeded! Model access group={group.name} Current cost: {_MAX_BUDGET}, Max budget: {_MAX_BUDGET}",
|
||||
)
|
||||
|
|
@ -8,13 +8,17 @@ team-admin gate can prove parity.
|
|||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Callable, Iterator, Mapping
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass, replace
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import MappingProxyType
|
||||
from typing import Final, Literal, assert_never
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic import JsonValue
|
||||
|
||||
|
|
@ -29,6 +33,7 @@ from tests.integration._support.client import (
|
|||
team_admin_permissions,
|
||||
)
|
||||
from tests.integration._support.database import read_rows
|
||||
from tests.integration._support.wire import Reply, Request, Wire, wire_server
|
||||
|
||||
Caller = Literal["proxy_admin", "team_admin", "member", "other_team_admin", "outsider", "org_admin", "other_org_admin"]
|
||||
CALLERS: Final[tuple[Caller, ...]] = (
|
||||
|
|
@ -156,13 +161,13 @@ def _spend_rows(gateway: Gateway, team_id: str, since: datetime, until: datetime
|
|||
return rows
|
||||
|
||||
|
||||
def _team_model_body(s: TeamScenario, name: str) -> dict[str, JsonValue]:
|
||||
def _team_model_body(s: TeamScenario, name: str, api_base: str | None = None) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model_name": name,
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_key": "integration-provider-key",
|
||||
"api_base": f"{s.gateway.upstream_url}/v1",
|
||||
"api_base": api_base or f"{s.gateway.upstream_url}/v1",
|
||||
},
|
||||
"model_info": {"team_id": s.team_id},
|
||||
}
|
||||
|
|
@ -440,8 +445,378 @@ def test_status_code(shared: TeamScenario, org_team: TeamScenario, route: Route,
|
|||
f"{caller} {call.method} {call.path}: {response.status_code} {response.text}"
|
||||
)
|
||||
if route.name in _BUDGET_ROUTES:
|
||||
assert read_rows(
|
||||
'SELECT max_budget FROM "LiteLLM_TeamTable" WHERE team_id = %s', (s.team_id,)
|
||||
) == [{"max_budget": _BUDGET_ROUTES[route.name] if response.status_code == 200 else _BUDGET_BEFORE}]
|
||||
assert read_rows('SELECT max_budget FROM "LiteLLM_TeamTable" WHERE team_id = %s', (s.team_id,)) == [
|
||||
{"max_budget": _BUDGET_ROUTES[route.name] if response.status_code == 200 else _BUDGET_BEFORE}
|
||||
]
|
||||
if response.status_code == 200 and route.cleanup is not None:
|
||||
route.cleanup(s, object_value(response.json()))
|
||||
|
||||
|
||||
ADMIN_ONLY_CALLERS: Final[tuple[tuple[str, Caller], ...]] = (
|
||||
("shared", "team_admin"),
|
||||
("shared", "member"),
|
||||
("shared", "outsider"),
|
||||
("org_team", "org_admin"),
|
||||
)
|
||||
ADMIN_ONLY_IDS: Final = tuple(caller for _, caller in ADMIN_ONLY_CALLERS)
|
||||
|
||||
|
||||
def _admin_only_team(request: pytest.FixtureRequest, fixture: str, scenario: Scenario) -> TeamScenario:
|
||||
team: Final = request.getfixturevalue(fixture)
|
||||
assert isinstance(team, TeamScenario)
|
||||
return replace(team, scenario=scenario)
|
||||
|
||||
|
||||
def _team_model(s: TeamScenario, api_base: str | None = None) -> tuple[str, str]:
|
||||
name: Final = f"matrix-{uuid.uuid4().hex}"
|
||||
created: Final = s.gateway.post("/model/new", _team_model_body(s, name, api_base))
|
||||
model_id: Final = string_value(object_value(created["model_info"])["id"])
|
||||
s.scenario.cleanups.callback(_delete_model_if_present, s.gateway, model_id)
|
||||
return model_id, name
|
||||
|
||||
|
||||
def _blocked(model_id: str) -> list[dict[str, object]]:
|
||||
return read_rows('SELECT blocked FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (model_id,))
|
||||
|
||||
|
||||
def _serve(s: TeamScenario, name: str, key: str) -> httpx.Response:
|
||||
return s.gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": name, "messages": [{"role": "user", "content": "serving state"}]},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
def _is_model_discovery(request: Request) -> bool:
|
||||
return (request.method, request.target) == ("GET", "/v1/models")
|
||||
|
||||
|
||||
def _model_serving_reply(request: Request) -> Reply:
|
||||
if _is_model_discovery(request):
|
||||
return Reply(body=b'{"object":"list","data":[]}')
|
||||
assert (request.method, request.target) == ("POST", "/v1/chat/completions"), request
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-model-blocked-control",
|
||||
"object": "chat.completion",
|
||||
"created": 0,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "model serving control"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _provider_calls(wire: Wire) -> tuple[Request, ...]:
|
||||
return tuple(entry for entry in wire.drain() if not _is_model_discovery(entry))
|
||||
|
||||
|
||||
def _assert_serving_state(s: TeamScenario, name: str, wire: Wire, *, blocked: bool) -> None:
|
||||
served: Final = _serve(s, name, s.keys["team_admin"])
|
||||
if blocked:
|
||||
assert (served.status_code, served.json()) == _MODEL_BLOCKED_REFUSED, served.text
|
||||
assert _provider_calls(wire) == (), f"Blocked model reached the provider: {served.text}"
|
||||
return
|
||||
|
||||
assert served.status_code == 200, served.text
|
||||
assert served.json()["choices"][0]["message"]["content"] == "model serving control", served.text
|
||||
received: Final = _provider_calls(wire)
|
||||
assert [
|
||||
(entry.method, entry.target, entry.headers["authorization"], json.loads(entry.body)) for entry in received
|
||||
] == [
|
||||
(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
"Bearer integration-provider-key",
|
||||
{"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "serving state"}]},
|
||||
)
|
||||
], served.text
|
||||
|
||||
|
||||
def _refusal(response: httpx.Response) -> JsonValue:
|
||||
"""The error body with the caller's masked user id elided, since it differs per caller."""
|
||||
body: Final = response.json()
|
||||
error: Final = object_value(body).get("error")
|
||||
if isinstance(error, dict) and isinstance(error.get("message"), str):
|
||||
return {
|
||||
**body,
|
||||
"error": {**error, "message": re.sub(r"Your user_id=\S+", "Your user_id=<caller>", error["message"])},
|
||||
}
|
||||
return body
|
||||
|
||||
|
||||
def _admin_only_route(route: str) -> tuple[int, JsonValue]:
|
||||
return 401, {
|
||||
"error": {
|
||||
"message": "Authentication Error, Only proxy admin can be used to generate, delete, update info for new "
|
||||
f"keys/users/teams. Route={route}. Your role=internal_user. Your user_id=<caller>",
|
||||
"type": "auth_error",
|
||||
"param": "None",
|
||||
"code": "401",
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
_BLOCK_REFUSALS: Final[Mapping[tuple[str, str], tuple[int, JsonValue]]] = MappingProxyType(
|
||||
{
|
||||
("POST", "/model/block"): _admin_only_route("/model/block"),
|
||||
("POST", "/model/unblock"): _admin_only_route("/model/unblock"),
|
||||
}
|
||||
)
|
||||
_PATCH_BLOCKED_REFUSED: Final[tuple[int, JsonValue]] = (
|
||||
403,
|
||||
{
|
||||
"error": {
|
||||
"message": "Only proxy admins can change a model's blocked flag.",
|
||||
"type": "auth_error",
|
||||
"param": "blocked",
|
||||
"code": "403",
|
||||
}
|
||||
},
|
||||
)
|
||||
_MODEL_BLOCKED_REFUSED: Final[tuple[int, JsonValue]] = (
|
||||
403,
|
||||
{
|
||||
"error": {
|
||||
"message": "litellm.PermissionDeniedError: Model is blocked",
|
||||
"type": "permission_error",
|
||||
"param": None,
|
||||
"code": "403",
|
||||
}
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("fixture", "caller"), ADMIN_ONLY_CALLERS, ids=ADMIN_ONLY_IDS)
|
||||
def test_only_proxy_admin_flips_a_models_blocked_flag(
|
||||
request: pytest.FixtureRequest, fixture: str, caller: Caller
|
||||
) -> None:
|
||||
with (
|
||||
gateway_from_environment() as gateway,
|
||||
gateway.scenario() as scenario,
|
||||
wire_server(_model_serving_reply) as wire,
|
||||
):
|
||||
s: Final = _admin_only_team(request, fixture, scenario)
|
||||
model_id, name = _team_model(s, f"{wire.url}/v1")
|
||||
internal_key: Final = (
|
||||
s.scenario.key(user_id=s.scenario.user(user_role="internal_user")) if caller == "team_admin" else None
|
||||
)
|
||||
for blocked in (True, False):
|
||||
s.gateway.post("/model/block" if blocked else "/model/unblock", {"model_id": model_id})
|
||||
attempts = (
|
||||
Call("POST", "/model/unblock" if blocked else "/model/block", {"model_id": model_id}),
|
||||
Call("PATCH", f"/model/{model_id}/update", {"blocked": not blocked}),
|
||||
)
|
||||
for call in attempts:
|
||||
response = s.gateway.request(call.method, call.path, call.body, key=s.keys[caller])
|
||||
if call.method == "PATCH" and caller != "team_admin":
|
||||
assert response.status_code == 403, f"{caller} {call.method} {call.path}: {response.text}"
|
||||
else:
|
||||
expected = (
|
||||
_PATCH_BLOCKED_REFUSED if call.method == "PATCH" else _BLOCK_REFUSALS[call.method, call.path]
|
||||
)
|
||||
assert (response.status_code, _refusal(response)) == expected, (
|
||||
f"{caller} {call.method} {call.path}: {response.text}"
|
||||
)
|
||||
assert _blocked(model_id) == [{"blocked": blocked}], (
|
||||
f"{caller} {call.method} {call.path} flipped blocked"
|
||||
)
|
||||
if internal_key is not None:
|
||||
path = "/model/unblock" if blocked else "/model/block"
|
||||
internal_response = s.gateway.request("POST", path, {"model_id": model_id}, key=internal_key)
|
||||
assert (internal_response.status_code, _refusal(internal_response)) == _BLOCK_REFUSALS["POST", path], (
|
||||
internal_response.text
|
||||
)
|
||||
assert _blocked(model_id) == [{"blocked": blocked}], internal_response.text
|
||||
_assert_serving_state(s, name, wire, blocked=blocked)
|
||||
|
||||
|
||||
def test_a_key_allowed_the_block_routes_still_cannot_flip_a_models_blocked_flag(shared: TeamScenario) -> None:
|
||||
with wire_server(_model_serving_reply) as wire:
|
||||
user: Final = shared.scenario.user(user_role="internal_user")
|
||||
key: Final = shared.scenario.key(user_id=user, allowed_routes=["/model/block", "/model/unblock"])
|
||||
model_id, name = _team_model(shared, f"{wire.url}/v1")
|
||||
for blocked in (True, False):
|
||||
shared.gateway.post("/model/block" if blocked else "/model/unblock", {"model_id": model_id})
|
||||
call = Call("POST", "/model/unblock" if blocked else "/model/block", {"model_id": model_id})
|
||||
response = shared.gateway.request(call.method, call.path, call.body, key=key)
|
||||
assert (response.status_code, response.json()) == _PATCH_BLOCKED_REFUSED, response.text
|
||||
assert _blocked(model_id) == [{"blocked": blocked}], response.text
|
||||
_assert_serving_state(shared, name, wire, blocked=blocked)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _model_creation_disabled_for_internal_users(gateway: Gateway) -> Iterator[None]:
|
||||
setting: Final = "disable_model_add_for_internal_users"
|
||||
original: Final = object_value(gateway.get("/get/ui_settings")["values"]).get(setting, False)
|
||||
response: Final = gateway.request("PATCH", "/update/ui_settings", {setting: True})
|
||||
assert response.status_code == 200, response.text
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
restored: Final = gateway.request("PATCH", "/update/ui_settings", {setting: original})
|
||||
assert restored.status_code == 200, restored.text
|
||||
|
||||
|
||||
def _model_rows(name: str) -> list[dict[str, object]]:
|
||||
return read_rows(
|
||||
'SELECT model_id FROM "LiteLLM_ProxyModelTable" '
|
||||
"WHERE model_name = %s OR model_info->>'team_public_model_name' = %s",
|
||||
(name, name),
|
||||
)
|
||||
|
||||
|
||||
def _listed_names(gateway: Gateway, key: str) -> tuple[frozenset[str], frozenset[str]]:
|
||||
info: Final = gateway.get("/model/info")["data"]
|
||||
models: Final = gateway.request("GET", "/v1/models", key=key)
|
||||
assert models.status_code == 200, models.text
|
||||
listed: Final = models.json()["data"]
|
||||
assert isinstance(info, list) and isinstance(listed, list), models.text
|
||||
info_names: Final = frozenset(
|
||||
name
|
||||
for entry in map(object_value, info)
|
||||
for name in (entry["model_name"], object_value(entry["model_info"]).get("team_public_model_name"))
|
||||
if isinstance(name, str)
|
||||
)
|
||||
return info_names, frozenset(string_value(object_value(entry)["id"]) for entry in listed)
|
||||
|
||||
|
||||
def test_internal_user_model_creation_prohibition_refuses_team_admin_but_not_proxy_admin(
|
||||
shared: TeamScenario, peer: Gateway
|
||||
) -> None:
|
||||
with shared.gateway.scenario() as scenario:
|
||||
s: Final = replace(shared, scenario=scenario)
|
||||
team_admin: Final = s.keys["team_admin"]
|
||||
refused_name: Final = f"matrix-{uuid.uuid4().hex}"
|
||||
admin_name: Final = f"matrix-{uuid.uuid4().hex}"
|
||||
with _model_creation_disabled_for_internal_users(peer):
|
||||
refused: Final = s.gateway.request("POST", "/model/new", _team_model_body(s, refused_name), key=team_admin)
|
||||
assert (refused.status_code, refused.json()) == (
|
||||
403,
|
||||
{
|
||||
"error": {
|
||||
"message": "Model creation is disabled for internal users by disable_model_add_for_internal_users.",
|
||||
"type": "auth_error",
|
||||
"param": "None",
|
||||
"code": "403",
|
||||
}
|
||||
},
|
||||
), refused.text
|
||||
assert _model_rows(refused_name) == []
|
||||
info_names, listed = _listed_names(s.gateway, team_admin)
|
||||
assert refused_name not in info_names and refused_name not in listed, (info_names, listed)
|
||||
created: Final = s.gateway.post("/model/new", _team_model_body(s, admin_name))
|
||||
s.scenario.cleanups.callback(
|
||||
_delete_model_if_present, s.gateway, string_value(object_value(created["model_info"])["id"])
|
||||
)
|
||||
assert _model_rows(admin_name) == [{"model_id": object_value(created["model_info"])["id"]}]
|
||||
assert admin_name in _listed_names(s.gateway, team_admin)[1]
|
||||
allowed_name: Final = f"matrix-{uuid.uuid4().hex}"
|
||||
allowed: Final = s.gateway.request("POST", "/model/new", _team_model_body(s, allowed_name), key=team_admin)
|
||||
assert allowed.status_code == 200, allowed.text
|
||||
s.scenario.cleanups.callback(
|
||||
_delete_model_if_present, s.gateway, string_value(object_value(allowed.json()["model_info"])["id"])
|
||||
)
|
||||
assert len(_model_rows(allowed_name)) == 1
|
||||
|
||||
|
||||
def _credential_row(name: str) -> list[dict[str, object]]:
|
||||
return read_rows(
|
||||
'SELECT credential_name, credential_values, credential_info, updated_at::text AS updated_at FROM "LiteLLM_CredentialsTable" '
|
||||
"WHERE credential_name = %s",
|
||||
(name,),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("contains_slash", (False, True), ids=("plain-name", "slash-name"))
|
||||
@pytest.mark.parametrize(("fixture", "caller"), ADMIN_ONLY_CALLERS, ids=ADMIN_ONLY_IDS)
|
||||
def test_only_proxy_admin_reads_or_changes_credentials(
|
||||
request: pytest.FixtureRequest, fixture: str, caller: Caller, contains_slash: bool
|
||||
) -> None:
|
||||
secret: Final = f"synthetic-shared-credential-{uuid.uuid4().hex}"
|
||||
with gateway_from_environment() as gateway, gateway.scenario() as scenario:
|
||||
s: Final = _admin_only_team(request, fixture, scenario)
|
||||
name: Final = f"credential/{uuid.uuid4().hex}" if contains_slash else f"credential-{uuid.uuid4().hex}"
|
||||
attempted: Final = f"credential/{uuid.uuid4().hex}" if contains_slash else f"credential-{uuid.uuid4().hex}"
|
||||
s.gateway.post(
|
||||
"/credentials",
|
||||
{"credential_name": name, "credential_values": {"api_key": secret}, "credential_info": {"team": "shared"}},
|
||||
)
|
||||
s.scenario.cleanups.callback(s.gateway.request, "DELETE", f"/credentials/{name}")
|
||||
s.scenario.cleanups.callback(s.gateway.request, "DELETE", f"/credentials/{attempted}")
|
||||
model_name: Final = f"credential-model-{uuid.uuid4().hex}"
|
||||
created_model: Final = s.gateway.post(
|
||||
"/model/new",
|
||||
{
|
||||
"model_name": model_name,
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_base": f"{s.gateway.upstream_url}/v1",
|
||||
"litellm_credential_name": name,
|
||||
},
|
||||
},
|
||||
)
|
||||
model_id: Final = string_value(object_value(created_model["model_info"])["id"])
|
||||
s.scenario.cleanups.callback(_delete_model_if_present, s.gateway, model_id)
|
||||
credential_for_model: Final = s.gateway.request("GET", f"/credentials/by_model/{model_id}")
|
||||
assert credential_for_model.status_code == 200, credential_for_model.text
|
||||
assert credential_for_model.json() == {
|
||||
"credential_name": f"{model_name}-credential-{model_id}",
|
||||
"credential_values": {"api_base": f"{s.gateway.upstream_url}/v1"},
|
||||
"credential_info": {},
|
||||
}, credential_for_model.text
|
||||
stored: Final = _credential_row(name)
|
||||
assert len(stored) == 1, stored
|
||||
attempts: Final = (
|
||||
Call(
|
||||
"POST",
|
||||
"/credentials",
|
||||
{"credential_name": attempted, "credential_values": {"api_key": "k"}, "credential_info": {}},
|
||||
),
|
||||
Call(
|
||||
"PATCH",
|
||||
f"/credentials/{name}",
|
||||
{"credential_name": name, "credential_values": {"api_key": "overwritten"}, "credential_info": {}},
|
||||
),
|
||||
Call("DELETE", f"/credentials/{name}"),
|
||||
Call("GET", "/credentials"),
|
||||
Call("GET", f"/credentials/by_name/{name}"),
|
||||
Call("GET", f"/credentials/by_model/{model_id}"),
|
||||
)
|
||||
for call in attempts:
|
||||
response = s.gateway.request(call.method, call.path, call.body, key=s.keys[caller])
|
||||
assert (response.status_code, _refusal(response)) == _admin_only_route(call.path), (
|
||||
f"{caller} {call.method} {call.path}: {response.text}"
|
||||
)
|
||||
assert secret not in response.text, f"{caller} {call.method} {call.path} exposed the credential"
|
||||
assert _credential_row(name) == stored, f"{caller} {call.method} {call.path} changed the credential row"
|
||||
assert _credential_row(attempted) == [], f"{caller} {call.method} {call.path} wrote a credential"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("fixture", "caller"),
|
||||
(("shared", "member"), ("shared", "outsider"), ("org_team", "org_admin")),
|
||||
ids=("member", "outsider", "org_admin"),
|
||||
)
|
||||
def test_non_team_admin_patching_the_blocked_flag_gets_the_blocked_flag_refusal(
|
||||
request: pytest.FixtureRequest, fixture: str, caller: Caller
|
||||
) -> None:
|
||||
pytest.skip(
|
||||
'BUG: PATCH /model/{id}/update {"blocked": ...} by a non-team-admin returns the auto-router refusal '
|
||||
"instead of the blocked-flag refusal (https://github.com/BerriAI/litellm/issues/44792)"
|
||||
)
|
||||
with gateway_from_environment() as gateway, gateway.scenario() as scenario:
|
||||
s: Final = _admin_only_team(request, fixture, scenario)
|
||||
model_id, _ = _team_model(s)
|
||||
response: Final = s.gateway.request("PATCH", f"/model/{model_id}/update", {"blocked": True}, key=s.keys[caller])
|
||||
assert (response.status_code, _refusal(response)) == _PATCH_BLOCKED_REFUSED, response.text
|
||||
assert _blocked(model_id) == [{"blocked": False}], response.text
|
||||
|
|
|
|||
|
|
@ -1,11 +1,14 @@
|
|||
import json
|
||||
import uuid
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic import BaseModel, JsonValue
|
||||
|
||||
from tests.integration._support.client import Gateway, object_value, string_value
|
||||
from tests.integration._support.database import read_rows
|
||||
from tests.integration._support.wire import Reply, Request, Wire, wire_server
|
||||
|
||||
|
||||
def model_identity(gateway: Gateway, alias: str) -> str:
|
||||
|
|
@ -24,11 +27,12 @@ def test_model_block_changes_actual_route_and_leaves_other_route_working(gateway
|
|||
gateway.chat(model)
|
||||
gateway.chat(other)
|
||||
gateway.post("/model/block", {"model_id": identity})
|
||||
assert read_rows(
|
||||
'SELECT blocked FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)
|
||||
) == [{"blocked": True}]
|
||||
assert read_rows('SELECT blocked FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == [
|
||||
{"blocked": True}
|
||||
]
|
||||
response: Final = gateway.request(
|
||||
"POST", "/v1/chat/completions",
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": "blocked deployment"}]},
|
||||
)
|
||||
assert response.status_code == 403, response.text
|
||||
|
|
@ -36,15 +40,18 @@ def test_model_block_changes_actual_route_and_leaves_other_route_working(gateway
|
|||
assert response.json()["error"]["message"] == "litellm.PermissionDeniedError: Model is blocked"
|
||||
assert object_value(gateway.chat(other)["usage"])["total_tokens"] == 40
|
||||
gateway.post("/model/unblock", {"model_id": identity})
|
||||
assert read_rows(
|
||||
'SELECT blocked FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)
|
||||
) == [{"blocked": False}]
|
||||
assert read_rows('SELECT blocked FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == [
|
||||
{"blocked": False}
|
||||
]
|
||||
assert object_value(gateway.chat(model)["usage"])["total_tokens"] == 40
|
||||
|
||||
|
||||
@pytest.mark.covers("mgmt.router_settings.update.changes_observed_attempt_count")
|
||||
def test_saved_retry_setting_controls_real_attempts_and_restores(gateway: Gateway) -> None:
|
||||
with httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream, gateway.scenario() as scenario:
|
||||
with (
|
||||
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
original: Final = object_value(gateway.get("/router/settings")["current_values"])["num_retries"]
|
||||
provider_model: Final = f"retry-{uuid.uuid4().hex}"
|
||||
model: Final = scenario.model(model=f"openai/{provider_model}", input_cost_per_token=0, output_cost_per_token=0)
|
||||
|
|
@ -63,8 +70,12 @@ def test_saved_retry_setting_controls_real_attempts_and_restores(gateway: Gatewa
|
|||
assert configured.status_code == 200, configured.text
|
||||
upstream.get("/__observations").raise_for_status()
|
||||
response: Final = gateway.request(
|
||||
"POST", "/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": f"{provider_model} attempt {generation}"}]},
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"messages": [{"role": "user", "content": f"{provider_model} attempt {generation}"}],
|
||||
},
|
||||
)
|
||||
observed: Final = upstream.get("/__observations")
|
||||
observed.raise_for_status()
|
||||
|
|
@ -82,28 +93,210 @@ def test_saved_retry_setting_controls_real_attempts_and_restores(gateway: Gatewa
|
|||
assert object_value(gateway.get("/router/settings")["current_values"])["num_retries"] == original
|
||||
|
||||
|
||||
class _MaskedCredential(BaseModel):
|
||||
credential_name: str
|
||||
credential_info: dict[str, JsonValue]
|
||||
credential_values: dict[str, JsonValue]
|
||||
|
||||
|
||||
def _openai_reply(request: Request) -> Reply:
|
||||
assert (request.method, request.target) == ("POST", "/v1/chat/completions"), request
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-credential-merge",
|
||||
"object": "chat.completion",
|
||||
"created": 0,
|
||||
"model": "gpt-4o-mini",
|
||||
"choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}],
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _stored_credential_values(name: str) -> dict[str, object]:
|
||||
rows: Final = read_rows(
|
||||
'SELECT credential_values FROM "LiteLLM_CredentialsTable" WHERE credential_name = %s', (name,)
|
||||
)
|
||||
assert len(rows) == 1, rows
|
||||
return object_value(rows[0]["credential_values"])
|
||||
|
||||
|
||||
def _masked_credential(gateway: Gateway, name: str) -> _MaskedCredential:
|
||||
response: Final = gateway.request("GET", f"/credentials/by_name/{name}")
|
||||
assert response.status_code == 200, response.text
|
||||
return _MaskedCredential.model_validate_json(response.content)
|
||||
|
||||
|
||||
def _patch_credential(gateway: Gateway, name: str, body: dict[str, JsonValue]) -> None:
|
||||
response: Final = gateway.request("PATCH", f"/credentials/{name}", {"credential_name": name, **body})
|
||||
assert response.status_code == 200, response.text
|
||||
assert response.json() == {"success": True, "message": "Credential updated successfully"}, response.text
|
||||
|
||||
|
||||
def _provider_call(gateway: Gateway, model: str, text: str, *wires: Wire) -> tuple[tuple[Request, ...], ...]:
|
||||
for wire in wires:
|
||||
wire.drain()
|
||||
response: Final = gateway.request(
|
||||
"POST", "/v1/chat/completions", {"model": model, "messages": [{"role": "user", "content": text}]}
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
assert object_value(response.json()["usage"])["total_tokens"] == 2, response.text
|
||||
return tuple(wire.drain() for wire in wires)
|
||||
|
||||
|
||||
def _sent_headers(request: Request) -> tuple[str, str | None]:
|
||||
return request.headers["authorization"], request.headers.get("openai-organization")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("contains_slash", (False, True), ids=("plain-name", "slash-name"))
|
||||
def test_partial_credential_patch_merges_values_and_deletes_only_the_named_field(
|
||||
gateway: Gateway, contains_slash: bool
|
||||
) -> None:
|
||||
api_key: Final = f"synthetic-merge-{uuid.uuid4().hex}"
|
||||
with (
|
||||
wire_server(_openai_reply) as original,
|
||||
wire_server(_openai_reply) as moved,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
name: Final = (
|
||||
f"team-{uuid.uuid4().hex}/cred-{uuid.uuid4().hex}" if contains_slash else f"credential-{uuid.uuid4().hex}"
|
||||
)
|
||||
gateway.post(
|
||||
"/credentials",
|
||||
{
|
||||
"credential_name": name,
|
||||
"credential_values": {
|
||||
"api_key": api_key,
|
||||
"api_base": f"{original.url}/v1",
|
||||
"organization": "org-merge",
|
||||
},
|
||||
"credential_info": {},
|
||||
},
|
||||
)
|
||||
scenario.cleanups.callback(gateway.request, "DELETE", f"/credentials/{name}")
|
||||
model: Final = scenario.model(api_key=None, api_base=None, litellm_credential_name=name)
|
||||
first, untouched = _provider_call(gateway, model, "before patch", original, moved)
|
||||
assert [_sent_headers(request) for request in first] == [(f"Bearer {api_key}", "org-merge")]
|
||||
assert untouched == ()
|
||||
stored: Final = _stored_credential_values(name)
|
||||
assert sorted(stored) == ["api_base", "api_key", "organization"], stored
|
||||
|
||||
_patch_credential(gateway, name, {"credential_values": {"api_base": f"{moved.url}/v1"}, "credential_info": {}})
|
||||
merged: Final = _stored_credential_values(name)
|
||||
assert sorted(merged) == ["api_base", "api_key", "organization"], merged
|
||||
assert (merged["api_key"], merged["organization"]) == (stored["api_key"], stored["organization"]), merged
|
||||
assert merged["api_base"] != stored["api_base"], merged
|
||||
assert _masked_credential(gateway, name) == _MaskedCredential(
|
||||
credential_name=name,
|
||||
credential_info={},
|
||||
credential_values={
|
||||
"api_key": "sy****" + api_key[-2:],
|
||||
"api_base": f"{moved.url}/v1",
|
||||
"organization": "org-merge",
|
||||
},
|
||||
)
|
||||
old_base, new_base = _provider_call(gateway, model, "after api_base patch", original, moved)
|
||||
assert old_base == ()
|
||||
assert [(request.method, request.target) for request in new_base] == [("POST", "/v1/chat/completions")]
|
||||
assert [_sent_headers(request) for request in new_base] == [(f"Bearer {api_key}", "org-merge")]
|
||||
|
||||
_patch_credential(gateway, name, {"credential_values_to_delete": ["organization"], "credential_info": {}})
|
||||
pruned: Final = _stored_credential_values(name)
|
||||
assert pruned == {"api_key": merged["api_key"], "api_base": merged["api_base"]}, pruned
|
||||
assert _masked_credential(gateway, name) == _MaskedCredential(
|
||||
credential_name=name,
|
||||
credential_info={},
|
||||
credential_values={"api_key": "sy****" + api_key[-2:], "api_base": f"{moved.url}/v1"},
|
||||
)
|
||||
old_base_after_delete, after_delete = _provider_call(gateway, model, "after field delete", original, moved)
|
||||
assert old_base_after_delete == ()
|
||||
assert [_sent_headers(request) for request in after_delete] == [(f"Bearer {api_key}", None)]
|
||||
|
||||
|
||||
def test_dashboard_credential_edit_keeps_the_stripped_secret_and_round_trips_credential_info(gateway: Gateway) -> None:
|
||||
api_key: Final = f"synthetic-dashboard-{uuid.uuid4().hex}"
|
||||
info: Final[dict[str, JsonValue]] = {"custom_llm_provider": "openai"}
|
||||
with (
|
||||
wire_server(_openai_reply) as original,
|
||||
wire_server(_openai_reply) as moved,
|
||||
gateway.scenario() as scenario,
|
||||
):
|
||||
name: Final = f"credential-{uuid.uuid4().hex}"
|
||||
gateway.post(
|
||||
"/credentials",
|
||||
{
|
||||
"credential_name": name,
|
||||
"credential_values": {"api_key": api_key, "api_base": f"{original.url}/v1", "organization": "org-dash"},
|
||||
"credential_info": info,
|
||||
},
|
||||
)
|
||||
scenario.cleanups.callback(gateway.request, "DELETE", f"/credentials/{name}")
|
||||
model: Final = scenario.model(api_key=None, api_base=None, litellm_credential_name=name)
|
||||
stored: Final = _stored_credential_values(name)
|
||||
|
||||
_patch_credential(
|
||||
gateway,
|
||||
name,
|
||||
{
|
||||
"credential_values": {"api_base": f"{moved.url}/v1", "organization": "org-dash"},
|
||||
"credential_info": info,
|
||||
},
|
||||
)
|
||||
edited: Final = _stored_credential_values(name)
|
||||
assert sorted(edited) == ["api_base", "api_key", "organization"], edited
|
||||
assert edited["api_key"] == stored["api_key"], edited
|
||||
assert read_rows(
|
||||
'SELECT credential_info FROM "LiteLLM_CredentialsTable" WHERE credential_name = %s', (name,)
|
||||
) == [{"credential_info": info}]
|
||||
assert _masked_credential(gateway, name) == _MaskedCredential(
|
||||
credential_name=name,
|
||||
credential_info=info,
|
||||
credential_values={
|
||||
"api_key": "sy****" + api_key[-2:],
|
||||
"api_base": f"{moved.url}/v1",
|
||||
"organization": "org-dash",
|
||||
},
|
||||
)
|
||||
old_base, new_base = _provider_call(gateway, model, "after dashboard edit", original, moved)
|
||||
assert old_base == ()
|
||||
assert [_sent_headers(request) for request in new_base] == [(f"Bearer {api_key}", "org-dash")]
|
||||
|
||||
|
||||
@pytest.mark.covers("mgmt.credential.update.saved_value_reaches_wire")
|
||||
def test_credential_value_update_and_model_reload_reach_provider(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario, httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream:
|
||||
with (
|
||||
gateway.scenario() as scenario,
|
||||
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
||||
):
|
||||
name: Final = f"credential-{uuid.uuid4().hex}"
|
||||
gateway.post("/credentials", {
|
||||
"credential_name": name, "credential_values": {"api_key": "synthetic-credential-first"}, "credential_info": {}
|
||||
})
|
||||
gateway.post(
|
||||
"/credentials",
|
||||
{
|
||||
"credential_name": name,
|
||||
"credential_values": {"api_key": "synthetic-credential-first"},
|
||||
"credential_info": {},
|
||||
},
|
||||
)
|
||||
|
||||
def remove_credential() -> None:
|
||||
response: Final = gateway.request("DELETE", f"/credentials/{name}")
|
||||
assert response.status_code == 200, response.text
|
||||
assert read_rows(
|
||||
'SELECT credential_name FROM "LiteLLM_CredentialsTable" WHERE credential_name = %s', (name,)
|
||||
) == []
|
||||
assert (
|
||||
read_rows('SELECT credential_name FROM "LiteLLM_CredentialsTable" WHERE credential_name = %s', (name,))
|
||||
== []
|
||||
)
|
||||
|
||||
scenario.cleanups.callback(remove_credential)
|
||||
model: Final = scenario.model(api_key=None, litellm_credential_name=name)
|
||||
identity: Final = model_identity(gateway, model)
|
||||
for value in ("synthetic-credential-first", "synthetic-credential-second"):
|
||||
patched: Final = gateway.request("PATCH", f"/credentials/{name}", {
|
||||
"credential_name": name, "credential_values": {"api_key": value}, "credential_info": {}
|
||||
})
|
||||
patched: Final = gateway.request(
|
||||
"PATCH",
|
||||
f"/credentials/{name}",
|
||||
{"credential_name": name, "credential_values": {"api_key": value}, "credential_info": {}},
|
||||
)
|
||||
assert patched.status_code == 200, patched.text
|
||||
rows: Final = read_rows(
|
||||
'SELECT credential_values FROM "LiteLLM_CredentialsTable" WHERE credential_name = %s', (name,)
|
||||
|
|
@ -113,10 +306,15 @@ def test_credential_value_update_and_model_reload_reach_provider(gateway: Gatewa
|
|||
assert isinstance(stored["api_key"], str) and stored["api_key"] != value
|
||||
for reload in (False, True):
|
||||
if reload:
|
||||
response: Final = gateway.request("PATCH", f"/model/{identity}/update", {"model_info": {"description": value}})
|
||||
response: Final = gateway.request(
|
||||
"PATCH", f"/model/{identity}/update", {"model_info": {"description": value}}
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
upstream.get("/__observations").raise_for_status()
|
||||
assert object_value(gateway.chat(model, text=f"{name} {value} reload={reload}")["usage"])["total_tokens"] == 40
|
||||
assert (
|
||||
object_value(gateway.chat(model, text=f"{name} {value} reload={reload}")["usage"])["total_tokens"]
|
||||
== 40
|
||||
)
|
||||
observed: Final = upstream.get("/__observations")
|
||||
observed.raise_for_status()
|
||||
assert len(observed.json()["requests"]) == 1, (value, reload, observed.text)
|
||||
|
|
|
|||
326
tests/integration/management/test_model_rename_propagation.py
Normal file
326
tests/integration/management/test_model_rename_propagation.py
Normal file
|
|
@ -0,0 +1,326 @@
|
|||
"""Renaming a database model carries every ``models`` allowlist that named it over to the new name.
|
||||
|
||||
Keys, teams, organizations, projects and users store public model names, and access groups store them as
|
||||
``access_model_names``. A caller restricted through any one of them must reach the renamed model under its
|
||||
new name, lose the old one, and see only the new name in ``/v1/models``.
|
||||
"""
|
||||
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
from dataclasses import dataclass
|
||||
from itertools import chain
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic import BaseModel, JsonValue, TypeAdapter
|
||||
|
||||
from integration._support.client import Gateway, Scenario, eventually, object_value, string_value
|
||||
from integration._support.database import read_rows
|
||||
from integration._support.wire import Reply, Request, Wire, wire_server
|
||||
|
||||
_PROVIDER_MODEL: Final = "gpt-4o-mini"
|
||||
_PROVIDER_API_KEY: Final = "integration-provider-key"
|
||||
_CALLER_KINDS: Final = ("key", "team", "organization", "project", "user", "access_group")
|
||||
_JSON_BODY: Final = TypeAdapter(dict[str, JsonValue])
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _Holder:
|
||||
kind: str
|
||||
table: str
|
||||
id_column: str
|
||||
object_id: str
|
||||
key: str
|
||||
listing_is_scoped: bool = True
|
||||
|
||||
|
||||
class _AccessGroup(BaseModel):
|
||||
access_group_id: str
|
||||
access_model_names: list[str]
|
||||
|
||||
|
||||
class _Usage(BaseModel):
|
||||
total_tokens: int
|
||||
|
||||
|
||||
class _Completion(BaseModel):
|
||||
model: str
|
||||
usage: _Usage
|
||||
|
||||
|
||||
class _Error(BaseModel):
|
||||
message: str
|
||||
type: str
|
||||
param: str | None
|
||||
code: str | None
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _RenameCall:
|
||||
kind: str
|
||||
status_code: int
|
||||
error: _Error
|
||||
requests: tuple[Request, ...]
|
||||
|
||||
|
||||
def _model_id(gateway: Gateway, name: str) -> str:
|
||||
entries: Final = gateway.get("/model/info")["data"]
|
||||
assert isinstance(entries, list)
|
||||
return next(
|
||||
string_value(object_value(entry["model_info"])["id"])
|
||||
for entry in map(object_value, entries)
|
||||
if entry["model_name"] == name
|
||||
)
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _access_group(gateway: Gateway, team_id: str, model: str) -> Iterator[str]:
|
||||
created: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/access_group",
|
||||
{
|
||||
"access_group_name": f"integration-{uuid.uuid4().hex}",
|
||||
"access_model_names": [model],
|
||||
"assigned_team_ids": [team_id],
|
||||
},
|
||||
)
|
||||
assert created.status_code == 201, created.text
|
||||
identity: Final = string_value(created.json()["access_group_id"])
|
||||
try:
|
||||
yield identity
|
||||
finally:
|
||||
deleted: Final = gateway.request("DELETE", f"/v1/access_group/{identity}")
|
||||
assert deleted.status_code == 204, deleted.text
|
||||
|
||||
|
||||
def _holders(scenario: Scenario, old: str) -> tuple[_Holder, ...]:
|
||||
team: Final = scenario.team(models=[old])
|
||||
organization: Final = scenario.organization(models=[old])
|
||||
org_team: Final = scenario.team(organization_id=organization, models=[old])
|
||||
project_team: Final = scenario.team()
|
||||
project: Final = scenario.project(project_team, models=[old])
|
||||
user: Final = scenario.user(user_role="internal_user", models=[old])
|
||||
key: Final = scenario.key(models=[old])
|
||||
return (
|
||||
_Holder("key", "LiteLLM_VerificationToken", "key_alias", "", key),
|
||||
_Holder("team", "LiteLLM_TeamTable", "team_id", team, scenario.key(team_id=team)),
|
||||
_Holder(
|
||||
"organization",
|
||||
"LiteLLM_OrganizationTable",
|
||||
"organization_id",
|
||||
organization,
|
||||
scenario.key(team_id=org_team, organization_id=organization),
|
||||
),
|
||||
_Holder(
|
||||
"project",
|
||||
"LiteLLM_ProjectTable",
|
||||
"project_id",
|
||||
project,
|
||||
scenario.key(team_id=project_team, project_id=project),
|
||||
listing_is_scoped=False,
|
||||
),
|
||||
_Holder("user", "LiteLLM_UserTable", "user_id", user, scenario.key(user_id=user), listing_is_scoped=False),
|
||||
)
|
||||
|
||||
|
||||
def _allowlist(holder: _Holder) -> list[dict[str, object]]:
|
||||
if holder.kind == "key":
|
||||
return read_rows(
|
||||
"SELECT models FROM \"LiteLLM_VerificationToken\" WHERE token = encode(sha256(%s::bytea), 'hex')",
|
||||
(holder.key,),
|
||||
)
|
||||
return read_rows(
|
||||
f'SELECT models FROM "{holder.table}" WHERE "{holder.id_column}" = %s', # noqa: S608 # table and column come from the fixed _holders tuple, never from input
|
||||
(holder.object_id,),
|
||||
)
|
||||
|
||||
|
||||
def _chat(gateway: Gateway, model: str, caller: _Holder) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "messages": [{"role": "user", "content": f"after rename {caller.kind}"}]},
|
||||
key=caller.key,
|
||||
)
|
||||
|
||||
|
||||
def _expected_provider_body(text: str) -> dict[str, JsonValue]:
|
||||
return {"model": _PROVIDER_MODEL, "messages": [{"role": "user", "content": text}]}
|
||||
|
||||
|
||||
def _is_model_discovery(request: Request) -> bool:
|
||||
return (request.method, request.target) == ("GET", "/v1/models")
|
||||
|
||||
|
||||
def _provider_calls(wire: Wire) -> tuple[Request, ...]:
|
||||
return tuple(entry for entry in wire.drain() if not _is_model_discovery(entry))
|
||||
|
||||
|
||||
def _rename_reply(request: Request) -> Reply:
|
||||
if _is_model_discovery(request):
|
||||
return Reply(body=json.dumps({"object": "list", "data": []}).encode())
|
||||
assert (request.method, request.target) == ("POST", "/v1/chat/completions"), request
|
||||
assert request.headers["authorization"] == f"Bearer {_PROVIDER_API_KEY}", request
|
||||
assert _JSON_BODY.validate_json(request.body) in tuple(
|
||||
_expected_provider_body(f"after rename {kind}") for kind in _CALLER_KINDS
|
||||
), request.body
|
||||
return Reply(
|
||||
body=json.dumps(
|
||||
{
|
||||
"id": "chatcmpl-rename",
|
||||
"object": "chat.completion",
|
||||
"created": 1,
|
||||
"model": _PROVIDER_MODEL,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"message": {"role": "assistant", "content": "renamed"},
|
||||
"finish_reason": "stop",
|
||||
}
|
||||
],
|
||||
"usage": {"prompt_tokens": 20, "completion_tokens": 20, "total_tokens": 40},
|
||||
}
|
||||
).encode()
|
||||
)
|
||||
|
||||
|
||||
def _listed(gateway: Gateway, key: str) -> list[str]:
|
||||
response: Final = gateway.request("GET", "/v1/models", key=key)
|
||||
assert response.status_code == 200, response.text
|
||||
entries: Final = response.json()["data"]
|
||||
assert isinstance(entries, list), response.text
|
||||
return sorted(string_value(object_value(entry)["id"]) for entry in entries)
|
||||
|
||||
|
||||
def _assert_listing(gateway: Gateway, caller: _Holder, *, served: str, gone: str | None) -> None:
|
||||
def converged(listed: list[str]) -> bool:
|
||||
if caller.listing_is_scoped:
|
||||
return listed == [served]
|
||||
# Project and user allowlists gate calls but not /v1/models, which then lists every proxy model.
|
||||
return served in listed and gone not in listed
|
||||
|
||||
eventually(lambda: _listed(gateway, caller.key), converged)
|
||||
|
||||
|
||||
def _served(gateway: Gateway, model: str, caller: _Holder) -> httpx.Response:
|
||||
served: Final = eventually(
|
||||
lambda: _chat(gateway, model, caller), lambda r: r.status_code == 200, return_last_on_timeout=True
|
||||
)
|
||||
assert served.status_code == 200, f"{caller.kind}: {served.text}"
|
||||
return served
|
||||
|
||||
|
||||
def _baseline_call(gateway: Gateway, model: str, caller: _Holder, wire: Wire) -> tuple[Request, ...]:
|
||||
served: Final = _served(gateway, model, caller)
|
||||
assert _Completion.model_validate_json(served.content) == _Completion(model=model, usage=_Usage(total_tokens=40)), (
|
||||
f"{caller.kind}: {served.text}"
|
||||
)
|
||||
_assert_listing(gateway, caller, served=model, gone=None)
|
||||
requests: Final = _provider_calls(wire)
|
||||
assert len(requests) == 1, f"{caller.kind}: expected one provider request, received {requests!r}"
|
||||
return requests
|
||||
|
||||
|
||||
def _renamed_call(
|
||||
gateway: Gateway,
|
||||
old: str,
|
||||
new: str,
|
||||
caller: _Holder,
|
||||
wire: Wire,
|
||||
) -> _RenameCall:
|
||||
served: Final = _served(gateway, new, caller)
|
||||
assert _Completion.model_validate_json(served.content) == _Completion(model=new, usage=_Usage(total_tokens=40)), (
|
||||
f"{caller.kind}: {served.text}"
|
||||
)
|
||||
requests: Final = _provider_calls(wire)
|
||||
assert len(requests) == 1, f"{caller.kind}: expected one provider request, received {requests!r}"
|
||||
refused: Final = _chat(gateway, old, caller)
|
||||
assert refused.status_code == 403, f"{caller.kind}: {refused.text}"
|
||||
error: Final = _Error.model_validate(object_value(refused.json())["error"])
|
||||
assert _provider_calls(wire) == (), f"{caller.kind}: the old name reached the provider"
|
||||
_assert_listing(gateway, caller, served=new, gone=old)
|
||||
return _RenameCall(caller.kind, refused.status_code, error, requests)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("rename_style", ("patch", "legacy_post"))
|
||||
def test_renamed_model_is_reachable_by_its_new_name_through_every_allowlist_that_named_it(
|
||||
gateway: Gateway, rename_style: str
|
||||
) -> None:
|
||||
with wire_server(_rename_reply) as wire, gateway.scenario() as scenario:
|
||||
old: Final = scenario.model(api_base=f"{wire.url}/v1", api_key=_PROVIDER_API_KEY)
|
||||
model_id: Final = _model_id(gateway, old)
|
||||
holders: Final = _holders(scenario, old)
|
||||
access_team: Final = scenario.team(models=["no-default-models"])
|
||||
access_key: Final = scenario.key(team_id=access_team)
|
||||
with _access_group(gateway, access_team, old) as access_group_id:
|
||||
callers: Final = (*holders, _Holder("access_group", "", "", access_group_id, access_key))
|
||||
baseline_requests: Final = tuple(_baseline_call(gateway, old, caller, wire) for caller in callers)
|
||||
|
||||
new: Final = f"integration-renamed-{uuid.uuid4().hex}"
|
||||
renamed: Final = (
|
||||
gateway.request("PATCH", f"/model/{model_id}/update", {"model_name": new})
|
||||
if rename_style == "patch"
|
||||
else gateway.request(
|
||||
"POST",
|
||||
"/model/update",
|
||||
{"model_name": new, "model_info": {"id": model_id}, "litellm_params": {}},
|
||||
)
|
||||
)
|
||||
assert renamed.status_code == 200, renamed.text
|
||||
assert read_rows('SELECT model_name FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (model_id,)) == [
|
||||
{"model_name": new}
|
||||
]
|
||||
|
||||
results: Final = tuple(_renamed_call(gateway, old, new, caller, wire) for caller in callers)
|
||||
for holder in holders:
|
||||
assert _allowlist(holder) == [{"models": [new]}], holder.kind
|
||||
group: Final = gateway.request("GET", f"/v1/access_group/{access_group_id}")
|
||||
assert group.status_code == 200, group.text
|
||||
assert _AccessGroup.model_validate_json(group.content) == _AccessGroup(
|
||||
access_group_id=access_group_id, access_model_names=[new]
|
||||
), group.text
|
||||
|
||||
denied_message: Final = (
|
||||
f"The requested model '{old}' is not available for this API key, or the model name is invalid. "
|
||||
"Check the models available to you and try again."
|
||||
)
|
||||
expected_denials: Final = {
|
||||
"key": _Error(message=denied_message, type="key_model_access_denied", param="model", code="403"),
|
||||
"team": _Error(message=denied_message, type="team_model_access_denied", param="model", code="403"),
|
||||
"organization": _Error(
|
||||
message=denied_message, type="team_model_access_denied", param="model", code="403"
|
||||
),
|
||||
"project": _Error(message=denied_message, type="project_model_access_denied", param="model", code="403"),
|
||||
"user": _Error(message=denied_message, type="user_model_access_denied", param="model", code="403"),
|
||||
"access_group": _Error(
|
||||
message=denied_message, type="team_model_access_denied", param="model", code="403"
|
||||
),
|
||||
}
|
||||
assert tuple(result.status_code for result in results) == (403,) * len(callers)
|
||||
assert {result.kind: result.error for result in results} == expected_denials, json.dumps(
|
||||
{result.kind: result.error.model_dump() for result in results}
|
||||
)
|
||||
provider_requests: Final = tuple(
|
||||
chain.from_iterable(baseline_requests)
|
||||
) + tuple(chain.from_iterable(result.requests for result in results))
|
||||
expected_request: Final = (
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
f"Bearer {_PROVIDER_API_KEY}",
|
||||
)
|
||||
expected_requests: Final = tuple(
|
||||
(*expected_request, _expected_provider_body(f"after rename {caller.kind}"))
|
||||
for caller in callers
|
||||
) * 2
|
||||
assert tuple(
|
||||
(
|
||||
request.method,
|
||||
request.target,
|
||||
request.headers["authorization"],
|
||||
_JSON_BODY.validate_json(request.body),
|
||||
)
|
||||
for request in provider_requests
|
||||
) == expected_requests
|
||||
|
|
@ -1,6 +1,7 @@
|
|||
"""Stored-config slots: credentials the proxy holds in its env, config or database reach only their owner.
|
||||
|
||||
Slots: A1 (virtual key raw value), A2 (master key), B2 (deployment ``api_key`` via ``/model/new``),
|
||||
Slots: A1 (virtual key raw value), A2 (master key), B2 (deployment ``api_key`` via ``/model/new``, kept through a
|
||||
partial ``POST /model/update`` and rotated through ``PATCH /model/{id}/update``),
|
||||
B3 (``/credentials`` entry named by ``litellm_credential_name``), B4 (deployment
|
||||
``aws_secret_access_key``), B4v and B4t (Vertex service-account JSON and the access token minted
|
||||
for it), B5 (team ``model_config`` credential override), E1 (guardrail ``api_key`` from config),
|
||||
|
|
@ -43,6 +44,7 @@ from integration._support.client import Gateway, Scenario, eventually, object_va
|
|||
from integration._support.database import read_rows
|
||||
from integration._support.sigv4 import encoded_path, signature
|
||||
from integration._support.wire import Reply, Request, wire_server
|
||||
from pydantic import BaseModel, JsonValue
|
||||
from integration.security._canary import MARKER, Canary, canary, find_canary
|
||||
from integration.security._sinks import (
|
||||
CONFIG_MODEL,
|
||||
|
|
@ -233,6 +235,27 @@ def _bearer(rig: Rig, marker: Canary, secret: Canary) -> None:
|
|||
)
|
||||
|
||||
|
||||
def _provider_chat_matches(rig: Rig, marker: Canary, secret: Canary, text: str) -> bool:
|
||||
delivered: Final = rig.provider.carrying(marker.value)
|
||||
actual: Final = tuple(
|
||||
(
|
||||
entry.method,
|
||||
entry.target,
|
||||
entry.headers.get("authorization"),
|
||||
json.loads(entry.body),
|
||||
)
|
||||
for entry in delivered
|
||||
)
|
||||
return actual == (
|
||||
(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
f"Bearer {secret.value}",
|
||||
{"model": "gpt-4o-mini", "messages": [{"role": "user", "content": text}]},
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
@pytest.mark.parametrize("outcome", OUTCOMES)
|
||||
def test_virtual_key_raw_value_authenticates_and_is_stored_only_as_a_hash(
|
||||
|
|
@ -339,6 +362,141 @@ def test_model_api_key_added_through_the_api_reaches_only_the_provider(
|
|||
)
|
||||
|
||||
|
||||
def _stored_litellm_params(model_id: str) -> dict[str, object]:
|
||||
rows: Final = read_rows('SELECT litellm_params FROM "LiteLLM_ProxyModelTable" WHERE model_id=%s', (model_id,))
|
||||
assert len(rows) == 1, rows
|
||||
return object_value(rows[0]["litellm_params"])
|
||||
|
||||
|
||||
class _ServedParams(BaseModel):
|
||||
model_name: str
|
||||
litellm_params: dict[str, JsonValue]
|
||||
|
||||
|
||||
def _served_params(gateway: Gateway, model_id: str) -> _ServedParams:
|
||||
response: Final = gateway.request("GET", "/model/info", params={"litellm_model_id": model_id})
|
||||
assert response.status_code == 200, response.text
|
||||
entries: Final = response.json()["data"]
|
||||
assert isinstance(entries, list) and len(entries) == 1, response.text
|
||||
return _ServedParams.model_validate(entries[0])
|
||||
|
||||
|
||||
def _served_openai_params(api_base: str, **changed: JsonValue) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_base": api_base,
|
||||
"allow_client_keepalive_override": False,
|
||||
"merge_reasoning_content_in_choices": False,
|
||||
"use_in_pass_through": False,
|
||||
"use_litellm_proxy": False,
|
||||
"use_xai_oauth": False,
|
||||
**changed,
|
||||
}
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
@pytest.mark.parametrize(
|
||||
("partial_params", "expected_params"),
|
||||
(
|
||||
pytest.param({"rpm": 10}, {"rpm": 10}, id="rpm"),
|
||||
pytest.param({"timeout": 30}, {"timeout": 30.0}, id="timeout"),
|
||||
),
|
||||
)
|
||||
def test_partial_legacy_model_update_keeps_the_stored_api_key_usable(
|
||||
rig: Rig, partial_params: dict[str, JsonValue], expected_params: dict[str, JsonValue]
|
||||
) -> None:
|
||||
b2: Final = canary("B2")
|
||||
marker: Final = canary(MARKER)
|
||||
api_base: Final = rig.provider.url + "/v1"
|
||||
with rig.proxy.scenario() as scenario:
|
||||
model, model_id = _model(scenario, {"model": "openai/gpt-4o-mini", "api_base": api_base, "api_key": b2.value})
|
||||
stored_before: Final = _stored_litellm_params(model_id)
|
||||
updated: Final = rig.proxy.request(
|
||||
"POST",
|
||||
"/model/update",
|
||||
{"model_name": model, "model_info": {"id": model_id}, "litellm_params": partial_params},
|
||||
)
|
||||
assert updated.status_code == 200, updated.text
|
||||
assert find_canary(updated.text, (b2,)) == (), f"/model/update echoed the stored key: {updated.text}"
|
||||
stored_after: Final = _stored_litellm_params(model_id)
|
||||
assert stored_after == {**stored_before, **expected_params}, (
|
||||
"A partial update rewrote stored params it was not sent, the api_key ciphertext included"
|
||||
)
|
||||
assert tuple((key, type(stored_after[key]), stored_after[key]) for key in expected_params) == tuple(
|
||||
(key, type(value), value) for key, value in expected_params.items()
|
||||
), stored_after
|
||||
_assert_stored_without_canary(
|
||||
"""SELECT litellm_params->>'api_key' FROM "LiteLLM_ProxyModelTable" WHERE model_id=%s""", (model_id,), b2
|
||||
)
|
||||
served_params: Final = _served_params(rig.proxy, model_id)
|
||||
assert served_params == _ServedParams(
|
||||
model_name=model, litellm_params=_served_openai_params(api_base, **expected_params)
|
||||
)
|
||||
assert tuple(
|
||||
(key, type(served_params.litellm_params[key]), served_params.litellm_params[key])
|
||||
for key in expected_params
|
||||
) == tuple((key, type(value), value) for key, value in expected_params.items()), served_params
|
||||
caller: Final = _caller(scenario, models=[model])
|
||||
_chat(rig.proxy, caller.key, model, "B2", marker, "success")
|
||||
assert _provider_chat_matches(rig, marker, b2, f"slot B2 {marker.value}"), (
|
||||
"The partial legacy update did not preserve the exact provider request"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
def test_patched_model_api_key_rotates_the_provider_bearer_and_is_never_stored_or_echoed(
|
||||
rig: Rig, request: pytest.FixtureRequest
|
||||
) -> None:
|
||||
started: Final = datetime.now(UTC)
|
||||
original: Final = canary("B2")
|
||||
rotated: Final = canary("B2")
|
||||
marker: Final = canary(MARKER)
|
||||
api_base: Final = rig.provider.url + "/v1"
|
||||
with rig.proxy.scenario() as scenario:
|
||||
model, model_id = _model(
|
||||
scenario, {"model": "openai/gpt-4o-mini", "api_base": api_base, "api_key": original.value}
|
||||
)
|
||||
stored_before: Final = _stored_litellm_params(model_id)
|
||||
patched: Final = rig.proxy.request(
|
||||
"PATCH", f"/model/{model_id}/update", {"litellm_params": {"api_key": rotated.value}}
|
||||
)
|
||||
assert patched.status_code == 200, patched.text
|
||||
assert find_canary(patched.text, (original, rotated)) == (), f"PATCH echoed a provider key: {patched.text}"
|
||||
_assert_stored_without_canary(
|
||||
"""SELECT litellm_params->>'api_key' FROM "LiteLLM_ProxyModelTable" WHERE model_id=%s""",
|
||||
(model_id,),
|
||||
rotated,
|
||||
)
|
||||
stored_after: Final = _stored_litellm_params(model_id)
|
||||
assert stored_after["api_key"] != stored_before["api_key"], "The rotated key was never written"
|
||||
assert {k: v for k, v in stored_after.items() if k != "api_key"} == {
|
||||
k: v for k, v in stored_before.items() if k != "api_key"
|
||||
}
|
||||
assert _served_params(rig.proxy, model_id) == _ServedParams(
|
||||
model_name=model, litellm_params=_served_openai_params(api_base)
|
||||
)
|
||||
caller: Final = _caller(scenario, models=[model])
|
||||
response, request_id = _chat(rig.proxy, caller.key, model, "B2", marker, "success")
|
||||
assert _provider_chat_matches(rig, marker, rotated, f"slot B2 {marker.value}"), (
|
||||
"The rotated key did not produce the exact provider request"
|
||||
)
|
||||
_finish(
|
||||
rig,
|
||||
rig.proxy,
|
||||
request,
|
||||
secrets=(original, rotated),
|
||||
marker=marker,
|
||||
response=response,
|
||||
request_id=request_id,
|
||||
caller=caller,
|
||||
ids={"model_id": model_id, "model": model},
|
||||
detail_routes=(f"/credentials/by_model/{model_id}",),
|
||||
reads=_reads(rig.proxy, {"/model/info": {"litellm_model_id": model_id}}),
|
||||
context="slot B2 rotated through PATCH",
|
||||
since=started,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.timeout(240)
|
||||
@pytest.mark.parametrize("outcome", OUTCOMES)
|
||||
def test_named_credential_reaches_only_the_provider(rig: Rig, outcome: str, request: pytest.FixtureRequest) -> None:
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue