From 50292cda2d0449cfd3c7ed05e564fd90b7b01879 Mon Sep 17 00:00:00 2001 From: Yash-Chindam Date: Sat, 3 Oct 2026 19:18:37 +0530 Subject: [PATCH] feat: restrict model classes, quotas, and privacy per tenant Section 7.1 makes resolving tenant quotas and permitted model classes a gateway responsibility and section 14 requires both to be restricted per tenant, but quota and model eligibility were global. Entitlements are governance and live in the catalog; the credential that binds a caller to a tenant stays in the environment (ROUTER_TENANT_KEYS). A tenant may restrict tiers or models, carry its own quota, set a privacy floor, deny external fallback, and raise the quality floor. An absent field defers to platform policy and never tightens an existing deployment. The entitlement is a hard filter applied before scoring. The privacy floor is applied before the cache lookup, so a tenant with a restricted floor cannot read an entry stored under public, and the route reason attributes the raise. Quota and cache are scoped to the tenant rather than the credential, so rotating a key resets neither. Co-Authored-By: Claude Opus 5.5 --- README.md | 34 +++++ config/registry.yaml | 30 ++++ src/llm_router/admission.py | 9 +- src/llm_router/app.py | 91 ++++++++++-- src/llm_router/config.py | 24 ++- src/llm_router/redis_state.py | 4 +- src/llm_router/registry.py | 80 +++++++++- src/llm_router/routing.py | 46 +++++- tests/integration/test_tenancy_api.py | 145 ++++++++++++++++++ tests/unit/test_tenancy.py | 203 ++++++++++++++++++++++++++ 10 files changed, 639 insertions(+), 27 deletions(-) create mode 100644 tests/integration/test_tenancy_api.py create mode 100644 tests/unit/test_tenancy.py diff --git a/README.md b/README.md index c13f339..a20b875 100644 --- a/README.md +++ b/README.md @@ -143,6 +143,40 @@ be served. A request can never introduce a model path, revision, or adapter. Send `routing.domain` to request a domain adapter; the router applies the promoted adapter with the largest measured quality gain for that base revision and task, or none at all. +## Tenants + +What a caller may use is governance and lives in the catalog; the credential that proves which +tenant is calling stays in the environment and is never committed. + +```bash +ROUTER_TENANT_KEYS="support-tooling:,clinical-research:" +``` + +Keys listed bare in `ROUTER_API_KEYS` belong to the `default` tenant, so an existing deployment +keeps working unchanged. Each tenant in [`config/registry.yaml`](config/registry.yaml) may set: + +| Field | Effect | +|---|---| +| `permitted_tiers` / `permitted_models` | Hard filter applied before scoring; it can never be outscored. Empty means unrestricted. | +| `quota_requests_per_minute` | The tenant's own limit; absent defers to the platform default. | +| `minimum_privacy` | A floor, never a ceiling. The request is raised to it before anything reads the class. | +| `allow_external_fallback` | `false` denies external routing outright; absent defers to platform policy. | +| `quality_floor` | Raises, and never lowers, the floor a request asked for. | + +An absent field never tightens an existing deployment: it defers to platform policy rather than +implying a restriction. + +Quota and cache are scoped to the tenant rather than the credential, so rotating a key neither +resets a quota nor orphans a cache, and two keys for one tenant draw on one quota. + +The privacy floor is applied before the cache lookup, not at routing. A tenant handling regulated +data cannot declare its traffic `public` and so become eligible for external routing or semantic +reuse, and it can never read an entry another tenant stored under `public`. When the class is +raised the route reason says so — the change is attributed, not silent. + +A request that no entitled model can serve is refused with `422` rather than downgraded to a +model never validated for the task. + ## Observability `GET /metrics` returns Prometheus exposition text and is intentionally unauthenticated so diff --git a/config/registry.yaml b/config/registry.yaml index 2d62ebc..03fea7c 100644 --- a/config/registry.yaml +++ b/config/registry.yaml @@ -179,3 +179,33 @@ deployments: model_checksums: small-specialist: sha256:mock-small-previous stage: deprecated + +# Per-tenant entitlement (sections 7.1 and 14). Entitlements are governance and +# belong here; the credentials that bind a caller to a tenant are supplied by +# ROUTER_TENANT_KEYS from the secret manager and are never committed. +tenants: + - id: default + description: Platform default for callers with no dedicated entitlement. + + - id: support-tooling + description: Ticket triage and field extraction only; no reasoning tier. + permitted_tiers: [small-specialist, general-local] + quota_requests_per_minute: 240 + minimum_privacy: private + + - id: clinical-research + description: Regulated data; never leaves the private inference environment. + permitted_tiers: [small-specialist, general-local, high-capability] + quota_requests_per_minute: 60 + # Raised to restricted no matter what the caller declares, so this tenant + # can never become eligible for external routing or semantic reuse. + minimum_privacy: restricted + quality_floor: 0.85 + # Stated explicitly rather than relying on the restricted floor alone. + allow_external_fallback: false + + - id: public-demo + description: Public demo traffic; may use the approved external fallback. + permitted_models: [small-specialist] + quota_requests_per_minute: 30 + allow_external_fallback: true diff --git a/src/llm_router/admission.py b/src/llm_router/admission.py index d881c28..e8ee735 100644 --- a/src/llm_router/admission.py +++ b/src/llm_router/admission.py @@ -36,13 +36,18 @@ def __init__(self, requests_per_minute: int) -> None: self._events: dict[str, deque[float]] = defaultdict(deque) self._lock = asyncio.Lock() - async def consume(self, subject: str, *, now: float | None = None) -> None: + async def consume( + self, subject: str, *, now: float | None = None, limit: int | None = None + ) -> None: + """Consume one request, against the subject's own limit when it has one.""" + timestamp = time.monotonic() if now is None else now cutoff = timestamp - 60 + effective_limit = self._limit if limit is None else limit async with self._lock: events = self._events[subject] while events and events[0] <= cutoff: events.popleft() - if len(events) >= self._limit: + if len(events) >= effective_limit: raise QuotaExceededError("request quota exceeded") events.append(timestamp) diff --git a/src/llm_router/app.py b/src/llm_router/app.py index 36fd1cd..5bfdc32 100644 --- a/src/llm_router/app.py +++ b/src/llm_router/app.py @@ -5,7 +5,9 @@ import uuid from collections.abc import AsyncIterator from contextlib import asynccontextmanager +from dataclasses import dataclass from pathlib import Path +from typing import Annotated import httpx from fastapi import Depends, FastAPI, Header, HTTPException, Request, Response, status @@ -49,10 +51,29 @@ ) from llm_router.observability import Metrics from llm_router.redis_state import RedisCacheStore, RedisFixedWindowQuota, RedisLike -from llm_router.registry import Registry, RegistryError, catalog_revisions, load_registry +from llm_router.registry import ( + Registry, + RegistryError, + TenantRecord, + catalog_revisions, + load_registry, + strictest_privacy, +) from llm_router.routing import NoEligibleModelError, Router, default_model_profiles +@dataclass(frozen=True) +class Principal: + """The authenticated caller: which tenant, and which credential proved it. + + Quota and cache scope use the tenant rather than the credential, so + rotating a key neither resets a tenant's quota nor orphans its cache. + """ + + tenant_id: str + credential_fingerprint: str + + def _redis_client(settings: Settings) -> RedisLike | None: """Build a shared-state client when a Redis URL is configured.""" @@ -162,7 +183,7 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: lifespan=lifespan, ) - async def authenticate(authorization: str | None = Header(default=None)) -> str: + async def authenticate(authorization: str | None = Header(default=None)) -> Principal: prefix = "Bearer " if authorization is None or not authorization.startswith(prefix): raise HTTPException( @@ -171,16 +192,25 @@ async def authenticate(authorization: str | None = Header(default=None)) -> str: headers={"WWW-Authenticate": "Bearer"}, ) token = authorization.removeprefix(prefix) - if not any( - secrets.compare_digest(token, candidate) - for candidate in runtime_settings.accepted_api_keys - ): + # Every candidate is compared so the work does not depend on which + # credential matched, and the match itself stays constant time. + matched: str | None = None + for candidate, tenant_id in runtime_settings.tenant_by_key.items(): + if secrets.compare_digest(token, candidate): + matched = tenant_id + if matched is None: raise HTTPException( status_code=status.HTTP_401_UNAUTHORIZED, detail="invalid bearer token", headers={"WWW-Authenticate": "Bearer"}, ) - return hashlib.sha256(token.encode()).hexdigest() + return Principal( + tenant_id=matched, + credential_fingerprint=hashlib.sha256(token.encode()).hexdigest(), + ) + + def _entitlement(principal: Principal) -> TenantRecord | None: + return catalog.tenant(principal.tenant_id) if catalog is not None else None @app.exception_handler(NoEligibleModelError) async def no_model_handler(_: Request, error: NoEligibleModelError) -> JSONResponse: @@ -409,12 +439,15 @@ async def _store_cache( scope = semantic_cache.scope(payload, tenant=tenant, model_revision=catalog_version) semantic_cache.store(scope, prompt, entry) - async def _consume_quota(subject: str) -> None: + async def _consume_quota(subject: str, tenant: TenantRecord | None) -> None: + # A tenant may carry its own limit; absent one the platform default + # applies, which is why None means "defer" rather than "unlimited". + limit = tenant.quota_requests_per_minute if tenant is not None else None if shared_quota is None: - await quota.consume(subject) + await quota.consume(subject, limit=limit) return window = int(time.time() // 60) - if not await shared_quota.consume(subject, window=window): + if not await shared_quota.consume(subject, window=window, limit=limit): raise QuotaExceededError("request quota exceeded") def _route_headers(decision: RouteDecision, cache_state: str) -> dict[str, str]: @@ -486,10 +519,31 @@ async def _stream_completion( async def chat_completions( payload: ChatCompletionRequest, response: Response, - subject: str = Depends(authenticate), + principal: Annotated[Principal, Depends(authenticate)], ) -> ChatCompletionResponse | StreamingResponse: started = time.perf_counter() - await _consume_quota(subject) + tenant = _entitlement(principal) + # Quota and cache are scoped to the tenant, not the credential, so + # rotating a key neither resets a quota nor orphans a cache. + subject = principal.tenant_id + await _consume_quota(subject, tenant) + + declared_privacy = payload.routing.privacy + if tenant is not None: + effective_privacy = strictest_privacy(declared_privacy, tenant.minimum_privacy) + if effective_privacy is not declared_privacy: + # Raised before anything reads the class, so cache eligibility, + # the cache key, and routing all agree on one value and a + # public-declared request can never reach restricted entries. + payload = payload.model_copy( + update={ + "routing": payload.routing.model_copy(update={"privacy": effective_privacy}) + } + ) + privacy_raised_from = ( + declared_privacy if payload.routing.privacy is not declared_privacy else None + ) + prompt = payload.prompt cache_key = build_cache_key( payload, tenant=subject, model_revision=catalog_version, prompt=prompt @@ -517,7 +571,18 @@ async def chat_completions( cached_task = decision_cache.get(prompt) if runtime_settings.cache_enabled else None telemetry.record_cache_event("router", "hit" if cached_task is not None else "miss") - decision = router.select(payload, task=cached_task) + decision = router.select( + payload, + task=cached_task, + permitted_models=( + catalog.permitted_models_for(tenant) if catalog is not None else None + ), + quality_floor=tenant.quality_floor if tenant is not None else 0.0, + tenant_allows_external=( + tenant.allow_external_fallback is not False if tenant is not None else True + ), + privacy_raised_from=privacy_raised_from, + ) if runtime_settings.cache_enabled: decision_cache.set(prompt, decision.task) telemetry.record_route(decision, privacy=payload.routing.privacy.value) diff --git a/src/llm_router/config.py b/src/llm_router/config.py index b4c02ba..7a1471a 100644 --- a/src/llm_router/config.py +++ b/src/llm_router/config.py @@ -12,6 +12,11 @@ class Settings(BaseSettings): environment: str = "development" api_keys: str = "dev-key" + # Binds credentials to tenant ids as "tenant:key,tenant:key". Entitlements + # are governance and live in the catalog; the secret that proves which + # tenant is calling stays in the environment and is never committed. + tenant_keys: str = "" + default_tenant: str = "default" max_concurrency: int = Field(default=32, ge=1) admission_timeout_seconds: float = Field(default=0.25, gt=0) quota_requests_per_minute: int = Field(default=120, ge=1) @@ -36,7 +41,24 @@ def reject_development_key_in_shared_environments(self) -> "Settings": @property def accepted_api_keys(self) -> frozenset[str]: - return frozenset(key.strip() for key in self.api_keys.split(",") if key.strip()) + return frozenset(self.tenant_by_key) + + @property + def tenant_by_key(self) -> dict[str, str]: + """Map every accepted credential to the tenant it authenticates. + + Keys listed without a tenant belong to the default tenant, so an + existing ROUTER_API_KEYS deployment keeps working unchanged. + """ + + bindings = { + key.strip(): self.default_tenant for key in self.api_keys.split(",") if key.strip() + } + for entry in self.tenant_keys.split(","): + tenant, separator, key = entry.partition(":") + if separator and tenant.strip() and key.strip(): + bindings[key.strip()] = tenant.strip() + return bindings @lru_cache diff --git a/src/llm_router/redis_state.py b/src/llm_router/redis_state.py index 9f6ea56..54429b6 100644 --- a/src/llm_router/redis_state.py +++ b/src/llm_router/redis_state.py @@ -64,11 +64,11 @@ def __init__(self, client: RedisLike, requests_per_minute: int, *, prefix: str = self._limit = requests_per_minute self._prefix = prefix - async def consume(self, subject: str, *, window: int) -> bool: + async def consume(self, subject: str, *, window: int, limit: int | None = None) -> bool: """Return whether the request fits inside the caller's quota.""" key = f"{self._prefix}{subject}:{window}" count = await self._client.incr(key) if count == 1: await self._client.expire(key, QUOTA_WINDOW_SECONDS * 2) - return count <= self._limit + return count <= (self._limit if limit is None else limit) diff --git a/src/llm_router/registry.py b/src/llm_router/registry.py index c1672f7..f027c4b 100644 --- a/src/llm_router/registry.py +++ b/src/llm_router/registry.py @@ -13,7 +13,21 @@ import yaml from pydantic import BaseModel, Field, model_validator -from llm_router.models import ModelProfile, TaskClass +from llm_router.models import ModelProfile, PrivacyClass, TaskClass + +# Privacy classes are ordered so a tenant floor can be compared with what a +# request declared; a floor may only ever raise the effective class. +PRIVACY_RANK: dict[PrivacyClass, int] = { + PrivacyClass.PUBLIC: 0, + PrivacyClass.PRIVATE: 1, + PrivacyClass.RESTRICTED: 2, +} + + +def strictest_privacy(left: PrivacyClass, right: PrivacyClass) -> PrivacyClass: + """Return whichever class is more sensitive.""" + + return left if PRIVACY_RANK[left] >= PRIVACY_RANK[right] else right class Quantization(StrEnum): @@ -148,6 +162,36 @@ class DeploymentRevision(BaseModel): previous_revision_id: str | None = None +class TenantRecord(BaseModel): + """What one caller is entitled to: which models, how much, how sensitive. + + Section 7.1 makes resolving tenant quotas and permitted model classes a + gateway responsibility, and section 14 requires both to be restricted per + tenant. Entitlements are governance, so they live in the catalog; the + credentials that bind a caller to a tenant stay in the environment and are + never committed here. + """ + + id: str + description: str = "" + # Empty means every servable tier or model, so a tenant that needs no + # restriction does not have to enumerate the catalog. + permitted_tiers: frozenset[ModelTier] = frozenset() + permitted_models: frozenset[str] = frozenset() + # None defers to the platform default rather than implying "unlimited". + quota_requests_per_minute: int | None = Field(default=None, ge=1) + # A floor, never a ceiling: a tenant handling regulated data must not be + # able to declare its traffic public and so become eligible for external + # routing or semantic reuse. Raising is always safe; lowering never happens. + minimum_privacy: PrivacyClass = PrivacyClass.PUBLIC + # None defers to platform policy, which already requires operator + # enablement and per-request opt-in; false denies external routing to this + # tenant outright. Absence never tightens an existing deployment, matching + # how an absent quota defers rather than meaning "none". + allow_external_fallback: bool | None = None + quality_floor: float = Field(default=0.0, ge=0.0, le=1.0) + + class RoutePolicy(BaseModel): """Operator-owned routing constraints applied before model scoring.""" @@ -171,6 +215,7 @@ class Registry(BaseModel): adapters: tuple[AdapterProfile, ...] = () benchmarks: tuple[BenchmarkRun, ...] = () deployments: tuple[DeploymentRevision, ...] = () + tenants: tuple[TenantRecord, ...] = () policy: RoutePolicy = RoutePolicy() @model_validator(mode="after") @@ -185,8 +230,41 @@ def validate_references(self) -> "Registry": f"adapter {adapter.id} references unknown base " f"{adapter.base_model_id}@{adapter.base_revision}" ) + tenant_ids = {tenant.id for tenant in self.tenants} + if len(tenant_ids) != len(self.tenants): + raise RegistryError("duplicate tenant identifiers in catalog") + for tenant in self.tenants: + unknown = tenant.permitted_models - model_ids + if unknown: + raise RegistryError( + f"tenant {tenant.id} permits unknown model(s): {', '.join(sorted(unknown))}" + ) return self + def tenant(self, tenant_id: str) -> TenantRecord | None: + return next((record for record in self.tenants if record.id == tenant_id), None) + + def permitted_models_for(self, tenant: TenantRecord | None) -> frozenset[str] | None: + """Resolve a tenant's entitlement to concrete model ids. + + Returns None when the tenant is unrestricted, which keeps the router's + hard filter free of tier lookups. An entitlement that resolves to + nothing is returned as an empty set so the request is refused rather + than silently widened to the whole catalog. + """ + + if tenant is None: + return None + if not tenant.permitted_tiers and not tenant.permitted_models: + return None + eligible = { + card.id + for card in self.servable_models() + if (not tenant.permitted_tiers or card.tier in tenant.permitted_tiers) + and (not tenant.permitted_models or card.id in tenant.permitted_models) + } + return frozenset(eligible) + def servable_models(self) -> tuple[ModelCard, ...]: return tuple( card diff --git a/src/llm_router/routing.py b/src/llm_router/routing.py index 2297d9e..ce60d0b 100644 --- a/src/llm_router/routing.py +++ b/src/llm_router/routing.py @@ -95,10 +95,22 @@ def classify_task(self, request: ChatCompletionRequest) -> TaskClass: ) def select( - self, request: ChatCompletionRequest, *, task: TaskClass | None = None + self, + request: ChatCompletionRequest, + *, + task: TaskClass | None = None, + permitted_models: frozenset[str] | None = None, + quality_floor: float = 0.0, + tenant_allows_external: bool = True, + privacy_raised_from: PrivacyClass | None = None, ) -> RouteDecision: task = task if task is not None else self.classify_task(request) estimated_tokens = max(1, len(request.prompt) // 4) + request.max_tokens + # The request arrives carrying its effective privacy class: the gateway + # raises it to the tenant floor before anything reads it, including the + # caches, so there is no declared-versus-effective split to get wrong. + effective_privacy = request.routing.privacy + floor = max(request.routing.quality_floor, quality_floor) candidates = [ profile @@ -106,9 +118,12 @@ def select( if profile.healthy and task in profile.supported_tasks and estimated_tokens <= profile.context_limit - and profile.quality >= request.routing.quality_floor - and self._privacy_allows(profile, request.routing.privacy) - and self._external_allows(profile, request) + and profile.quality >= floor + and self._privacy_allows(profile, effective_privacy) + and self._external_allows(profile, request, tenant_allows_external) + # A tenant entitlement is a hard restriction, like privacy: it is + # applied before scoring and can never be outscored. + and (permitted_models is None or profile.id in permitted_models) ] if request.model != "auto": @@ -117,7 +132,7 @@ def select( if not candidates: raise NoEligibleModelError( "no healthy model satisfies capability, context, quality, " - "privacy, and fallback policy" + "privacy, tenant entitlement, and fallback policy" ) def score(profile: ModelProfile) -> float: @@ -147,9 +162,15 @@ def score(profile: ModelProfile) -> float: ) reason = ( f"selected highest policy score among {len(candidates)} eligible model(s); " - f"task={task.value}, privacy={request.routing.privacy.value}, " + f"task={task.value}, privacy={effective_privacy.value}, " f"latency_tier={request.routing.latency_tier}" ) + if privacy_raised_from is not None: + # Raising the class is a policy decision and is attributed rather + # than applied silently. + reason += f"; raised from declared {privacy_raised_from.value} by the tenant floor" + if permitted_models is not None: + reason += f"; restricted to {len(permitted_models)} model(s) by tenant entitlement" if adapter is not None: reason += ( f"; applied adapter {adapter.id} for domain {adapter.domain} " @@ -169,7 +190,16 @@ def score(profile: ModelProfile) -> float: def _privacy_allows(profile: ModelProfile, privacy: PrivacyClass) -> bool: return profile.local or privacy == PrivacyClass.PUBLIC - def _external_allows(self, profile: ModelProfile, request: ChatCompletionRequest) -> bool: + def _external_allows( + self, + profile: ModelProfile, + request: ChatCompletionRequest, + tenant_allows_external: bool = True, + ) -> bool: + """External routing needs operator, tenant, and request agreement.""" + return profile.local or ( - self.external_fallback_enabled and request.routing.allow_external_fallback + self.external_fallback_enabled + and tenant_allows_external + and request.routing.allow_external_fallback ) diff --git a/tests/integration/test_tenancy_api.py b/tests/integration/test_tenancy_api.py new file mode 100644 index 0000000..906d727 --- /dev/null +++ b/tests/integration/test_tenancy_api.py @@ -0,0 +1,145 @@ +from fastapi.testclient import TestClient + +from llm_router.app import create_app +from llm_router.config import Settings +from llm_router.registry import Registry, load_registry + +CATALOG = load_registry("config/registry.yaml") +TENANT_KEYS = "support-tooling:support-key,clinical-research:clinical-key,public-demo:demo-key" + + +def build_client(*, registry: Registry | None = None, **overrides: object) -> TestClient: + settings = Settings( + api_keys="open-key", + tenant_keys=TENANT_KEYS, + **overrides, # type: ignore[arg-type] + ) + return TestClient(create_app(settings, registry=registry or CATALOG)) + + +def completion(client: TestClient, key: str, prompt: str, **routing: object) -> object: + return client.post( + "/v1/chat/completions", + headers={"Authorization": f"Bearer {key}"}, + json={ + "model": "auto", + "messages": [{"role": "user", "content": prompt}], + "routing": routing or {"privacy": "private"}, + }, + ) + + +def test_an_unknown_credential_is_still_rejected() -> None: + with build_client() as client: + assert completion(client, "not-a-key", "Classify this").status_code == 401 + + +def test_a_tenant_entitlement_narrows_the_models_considered() -> None: + with build_client() as client: + # support-tooling is entitled to the specialist and general tiers only. + response = completion(client, "support-key", "Summarize the ticket backlog") + + assert response.status_code == 200 + assert response.json()["model"] == "general-local" + assert "tenant entitlement" in response.headers["X-Route-Reason"] + + +def test_a_request_no_entitled_model_can_serve_is_refused_not_downgraded() -> None: + with build_client() as client: + # No model in support-tooling's tiers supports reasoning, so the request + # is refused explicitly rather than served by a model never validated + # for the task. + response = completion(client, "support-key", "Reason step by step about this proof") + + assert response.status_code == 422 + assert "tenant entitlement" in response.json()["error"]["message"] + + +def test_an_unrestricted_tenant_still_reaches_the_reasoning_tier() -> None: + with build_client() as client: + response = completion(client, "open-key", "Reason step by step about this proof") + + assert response.json()["model"] == "high-capability" + + +def test_a_tenant_floor_raises_a_declared_public_request_and_says_so() -> None: + with build_client() as client: + response = completion(client, "clinical-key", "Classify this ticket", privacy="public") + + assert response.status_code == 200 + assert "raised from declared public by the tenant floor" in response.headers["X-Route-Reason"] + + +def test_a_raised_floor_cannot_read_another_tenants_public_cache_entry() -> None: + prompt = "Classify this ticket as billing, technical, or other: card charged twice." + with build_client() as client: + first = completion(client, "open-key", prompt, privacy="public", temperature=0) + assert first.headers["X-Cache"] == "miss" + warm = completion(client, "open-key", prompt, privacy="public", temperature=0) + assert warm.headers["X-Cache"] == "exact" + + # Same prompt, same declared class, but this tenant's floor raises the + # request to restricted before any cache is consulted, so the entry + # stored under public must be unreachable. + isolated = completion(client, "clinical-key", prompt, privacy="public", temperature=0) + + assert isolated.headers["X-Cache"] == "miss" + + +def test_a_restricted_floor_tenant_is_never_cached_at_all() -> None: + prompt = "Summarize the patient intake notes" + with build_client() as client: + completion(client, "clinical-key", prompt, privacy="public", temperature=0) + repeat = completion(client, "clinical-key", prompt, privacy="public", temperature=0) + + # Restricted traffic is ineligible for exact reuse, and the floor applies + # even though the caller declared public. + assert repeat.headers["X-Cache"] == "miss" + + +def test_a_tenant_quota_is_enforced_from_the_catalog() -> None: + document = { + "models": [card.model_dump(mode="json") for card in CATALOG.models], + "tenants": [{"id": "public-demo", "quota_requests_per_minute": 1}], + } + with build_client(registry=Registry.model_validate(document)) as client: + assert completion(client, "demo-key", "Classify this").status_code == 200 + throttled = completion(client, "demo-key", "Classify this again") + + assert throttled.status_code == 429 + assert throttled.json()["error"]["type"] == "quota_exceeded" + + +def test_quota_is_shared_across_a_tenants_credentials_not_per_key() -> None: + document = { + "models": [card.model_dump(mode="json") for card in CATALOG.models], + "tenants": [{"id": "support-tooling", "quota_requests_per_minute": 1}], + } + settings = Settings( + api_keys="open-key", + tenant_keys="support-tooling:key-one,support-tooling:key-two", + ) + with TestClient(create_app(settings, registry=Registry.model_validate(document))) as client: + assert completion(client, "key-one", "Classify this").status_code == 200 + # A second credential for the same tenant draws on the same quota. + assert completion(client, "key-two", "Classify this").status_code == 429 + + +def test_a_tenant_without_a_quota_falls_back_to_the_platform_default() -> None: + with build_client(quota_requests_per_minute=1) as client: + assert completion(client, "open-key", "Classify this").status_code == 200 + assert completion(client, "open-key", "Classify this").status_code == 429 + + +def test_external_fallback_still_needs_the_operator_flag_even_for_a_permitted_tenant() -> None: + with build_client(external_fallback_enabled=False) as client: + response = completion( + client, + "demo-key", + "Classify this ticket", + privacy="public", + allow_external_fallback=True, + ) + + # public-demo is entitled to external routing, but the operator gate is shut. + assert response.json()["model"] == "small-specialist" diff --git a/tests/unit/test_tenancy.py b/tests/unit/test_tenancy.py new file mode 100644 index 0000000..1b0bc25 --- /dev/null +++ b/tests/unit/test_tenancy.py @@ -0,0 +1,203 @@ +import pytest + +from llm_router.admission import QuotaExceededError, SlidingWindowQuota +from llm_router.config import Settings +from llm_router.models import ChatCompletionRequest, PrivacyClass +from llm_router.registry import ( + ModelTier, + Registry, + RegistryError, + load_registry, + strictest_privacy, +) +from llm_router.routing import NoEligibleModelError, Router, default_model_profiles + +CATALOG = load_registry("config/registry.yaml") + + +def request_for(prompt: str, **routing: object) -> ChatCompletionRequest: + return ChatCompletionRequest.model_validate( + {"model": "auto", "messages": [{"role": "user", "content": prompt}], "routing": routing} + ) + + +def test_privacy_classes_are_ordered_by_sensitivity() -> None: + assert ( + strictest_privacy(PrivacyClass.PUBLIC, PrivacyClass.RESTRICTED) is PrivacyClass.RESTRICTED + ) + assert strictest_privacy(PrivacyClass.PRIVATE, PrivacyClass.PUBLIC) is PrivacyClass.PRIVATE + assert strictest_privacy(PrivacyClass.PRIVATE, PrivacyClass.PRIVATE) is PrivacyClass.PRIVATE + + +def test_committed_catalog_declares_tenant_entitlements() -> None: + tenant = CATALOG.tenant("clinical-research") + + assert tenant is not None + assert tenant.minimum_privacy is PrivacyClass.RESTRICTED + assert tenant.quota_requests_per_minute == 60 + assert tenant.allow_external_fallback is False + assert CATALOG.tenant("does-not-exist") is None + + +def test_an_unrestricted_tenant_resolves_to_no_restriction() -> None: + assert CATALOG.permitted_models_for(CATALOG.tenant("default")) is None + assert CATALOG.permitted_models_for(None) is None + + +def test_tier_entitlement_resolves_to_model_ids() -> None: + permitted = CATALOG.permitted_models_for(CATALOG.tenant("support-tooling")) + + assert permitted == frozenset({"small-specialist", "general-local"}) + + +def test_model_entitlement_narrows_further_than_the_tier() -> None: + permitted = CATALOG.permitted_models_for(CATALOG.tenant("public-demo")) + + assert permitted == frozenset({"small-specialist"}) + + +def test_an_entitlement_matching_nothing_refuses_rather_than_widening() -> None: + # One servable small-specialist card, and a tenant entitled only to the + # high-capability tier, so the entitlement resolves to nothing. An empty + # entitlement must stay empty rather than be read as "everything". + registry = Registry.model_validate( + { + "models": [ + card.model_dump(mode="json") + for card in CATALOG.models + if card.tier is ModelTier.SMALL_SPECIALIST + ], + "tenants": [{"id": "locked-out", "permitted_tiers": ["high-capability"]}], + } + ) + + assert registry.permitted_models_for(registry.tenant("locked-out")) == frozenset() + + +def test_a_tenant_cannot_permit_a_model_outside_the_catalog() -> None: + document = { + "models": [card.model_dump(mode="json") for card in CATALOG.models], + "tenants": [{"id": "bad", "permitted_models": ["model-that-does-not-exist"]}], + } + + with pytest.raises(RegistryError, match="permits unknown model"): + Registry.model_validate(document) + + +def test_duplicate_tenants_are_rejected() -> None: + document = { + "models": [card.model_dump(mode="json") for card in CATALOG.models], + "tenants": [{"id": "same"}, {"id": "same"}], + } + + with pytest.raises(RegistryError, match="duplicate tenant"): + Registry.model_validate(document) + + +def test_entitlement_is_a_hard_filter_that_scoring_cannot_outrank() -> None: + router = Router(profiles=default_model_profiles()) + + unrestricted = router.select(request_for("Reason carefully about this proof")) + restricted = router.select( + request_for("Reason carefully about this proof"), + permitted_models=frozenset({"general-local"}), + ) + + assert unrestricted.profile.id == "high-capability" + assert restricted.profile.id == "general-local" + assert "restricted to 1 model(s) by tenant entitlement" in restricted.reason + + +def test_an_entitlement_excluding_every_capable_model_is_refused() -> None: + router = Router(profiles=default_model_profiles()) + + with pytest.raises(NoEligibleModelError, match="tenant entitlement"): + router.select( + request_for("Reason carefully about this proof"), + # The small specialist supports neither reasoning nor this context. + permitted_models=frozenset({"small-specialist"}), + ) + + +def test_a_tenant_quality_floor_raises_but_never_lowers_the_request_floor() -> None: + router = Router(profiles=default_model_profiles()) + + raised = router.select(request_for("Classify this ticket"), quality_floor=0.95) + request_wins = router.select( + request_for("Classify this ticket", quality_floor=0.95), quality_floor=0.0 + ) + + assert raised.profile.quality >= 0.95 + assert request_wins.profile.quality >= 0.95 + + +def test_external_routing_needs_tenant_agreement_as_well() -> None: + router = Router(profiles=default_model_profiles(), external_fallback_enabled=True) + body = request_for("Reason about this", privacy="public", allow_external_fallback=True) + external_only = frozenset({"approved-external-fallback"}) + + permitted = router.select(body, permitted_models=external_only, tenant_allows_external=True) + + assert permitted.profile.id == "approved-external-fallback" + # Operator enablement and request opt-in are both already satisfied here, so + # the tenant gate alone is what decides between serving and refusing. + with pytest.raises(NoEligibleModelError, match="tenant entitlement"): + router.select(body, permitted_models=external_only, tenant_allows_external=False) + + +def test_a_local_model_still_wins_on_score_when_external_is_merely_eligible() -> None: + router = Router(profiles=default_model_profiles(), external_fallback_enabled=True) + body = request_for("Reason about this", privacy="public", allow_external_fallback=True) + + permitted = router.select(body, tenant_allows_external=True) + forbidden = router.select(body, tenant_allows_external=False) + + # Eligibility is not selection: the external model costs enough that the + # local high-capability tier outscores it either way. + assert permitted.profile.id == "high-capability" + assert forbidden.profile.id == "high-capability" + assert permitted.candidate_count == forbidden.candidate_count + 1 + + +def test_the_route_reason_attributes_a_raised_privacy_class() -> None: + router = Router(profiles=default_model_profiles()) + + decision = router.select( + request_for("Classify this ticket", privacy="restricted"), + privacy_raised_from=PrivacyClass.PUBLIC, + ) + + assert "raised from declared public by the tenant floor" in decision.reason + + +def test_credentials_bind_to_tenants_and_bare_keys_keep_working() -> None: + settings = Settings( + api_keys="legacy-key", + tenant_keys="support-tooling:support-key,clinical-research:clinical-key", + ) + + assert settings.tenant_by_key == { + "legacy-key": "default", + "support-key": "support-tooling", + "clinical-key": "clinical-research", + } + assert settings.accepted_api_keys == frozenset({"legacy-key", "support-key", "clinical-key"}) + + +def test_malformed_tenant_bindings_are_ignored_rather_than_trusted() -> None: + settings = Settings(api_keys="only-key", tenant_keys="no-colon-here,:,tenant:,:key") + + assert settings.tenant_by_key == {"only-key": "default"} + + +@pytest.mark.asyncio +async def test_quota_uses_the_tenant_limit_over_the_platform_default() -> None: + quota = SlidingWindowQuota(requests_per_minute=100) + + await quota.consume("tenant-a", now=1.0, limit=2) + await quota.consume("tenant-a", now=1.1, limit=2) + with pytest.raises(QuotaExceededError): + await quota.consume("tenant-a", now=1.2, limit=2) + + # A different tenant is accounted separately and may carry its own limit. + await quota.consume("tenant-b", now=1.3, limit=1)