"""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}", )