"""Lab 01 — the LLM gateway and model abstraction layer. The gateway is the platform's choke point, in the good sense: one place to enforce, one place to observe, one place to change providers. It owns normalization, routing, budget-aware fallback, caching, rate limiting, quotas, token accounting and cost attribution. Deterministic: the clock is injected, providers are scripted, money is integer micro-USD. pytest test_lab.py -v LAB_MODULE=solution pytest test_lab.py -v # the reference, must be green """ from __future__ import annotations import hashlib import math from dataclasses import dataclass, field, replace from enum import Enum from typing import Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple # ====================================================================================== # 1. The normalized request/response # ====================================================================================== class TaskClass(str, Enum): CHAT = "chat" EXTRACTION = "extraction" SUMMARIZATION = "summarization" CLASSIFICATION = "classification" REASONING = "reasoning" EMBEDDING = "embedding" @dataclass(frozen=True) class Message: role: str content: str @dataclass(frozen=True) class NormalizedRequest: """One request shape for every provider. The non-model fields — tenant, agent, classification, residency, latency budget, side_effecting — are why this type exists. A provider SDK's request object cannot carry them, and each one drives routing, caching or safety. """ messages: Tuple[Message, ...] task_class: TaskClass tenant: str agent_id: str max_output_tokens: int = 512 temperature: float = 0.0 data_classification: str = "internal" residency: str = "any" latency_budget_ms: int = 3000 side_effecting: bool = False cacheable: bool = True def __post_init__(self) -> None: # TODO: at least one message; max_output_tokens > 0; temperature in [0, 2]; # latency_budget_ms > 0. Otherwise ValueError. raise NotImplementedError def prompt_text(self) -> str: """``"role: content"`` per message, newline-joined.""" # TODO raise NotImplementedError def stable_prefix(self) -> str: """Everything BEFORE the last message — the cacheable part. Prefix caching rewards putting stable content first and volatile content last.""" # TODO raise NotImplementedError class FinishReason(str, Enum): STOP = "stop" LENGTH = "length" TOOL_CALL = "tool_call" CONTENT_FILTER = "content_filter" @dataclass(frozen=True) class Usage: input_tokens: int = 0 cached_input_tokens: int = 0 output_tokens: int = 0 def __post_init__(self) -> None: # TODO: no negatives; cached_input_tokens <= input_tokens raise NotImplementedError @property def total(self) -> int: return self.input_tokens + self.output_tokens @dataclass(frozen=True) class NormalizedResponse: text: str finish_reason: FinishReason usage: Usage deployment: str provider: str model: str latency_ms: int cost_micros: int = 0 cache: str = "miss" # "miss" | "exact" | "semantic" attempts: Tuple[str, ...] = () # ====================================================================================== # 2. Normalized errors — the hard half of an abstraction layer # ====================================================================================== class GatewayError(Exception): retryable = False fall_over = False # may we try the NEXT deployment? def __init__(self, message: str, *, provider: str = "") -> None: super().__init__(message) self.message = message self.provider = provider # TODO: set the two class flags on each subclass. # RateLimited, ProviderTimeout, ProviderUnavailable -> retryable AND fall_over # ContentFiltered -> NEITHER. Trying providers until one answers is "shopping for a # compliant model", which is exactly what a regulator asks about. # InvalidRequest, QuotaExceeded, NoRouteAvailable, BudgetExhausted -> neither. class RateLimited(GatewayError): """HTTP 429 from a provider.""" class ProviderTimeout(GatewayError): pass class ProviderUnavailable(GatewayError): """5xx.""" class ContentFiltered(GatewayError): """The provider's own safety system refused.""" class InvalidRequest(GatewayError): """A 4xx that is our fault; retrying sends the same bad request.""" class QuotaExceeded(GatewayError): pass class NoRouteAvailable(GatewayError): pass class BudgetExhausted(GatewayError): """The latency budget ran out before a fallback could be attempted.""" # ====================================================================================== # 3. Deployments and providers # ====================================================================================== class Capacity(str, Enum): PAYG = "payg" PROVISIONED = "provisioned" SELF_HOSTED = "self_hosted" @dataclass(frozen=True) class Deployment: """The routing target. A *model* is not one — the same model in two regions with two capacity types has two latencies, two prices and two residency answers.""" name: str provider: str model: str region: str capacity: Capacity input_micros_per_1k: int cached_input_micros_per_1k: int output_micros_per_1k: int expected_latency_ms: int max_classification: str = "restricted" weight: int = 1 def cost_micros(self, usage: Usage) -> int: """``((fresh*c_in + cached*c_cache + out*c_out) // 1000)``, dividing LAST.""" # TODO raise NotImplementedError ProviderAdapter = Callable[[NormalizedRequest, Deployment], NormalizedResponse] _CLASSIFICATION_ORDER = ("public", "internal", "confidential", "restricted") def classification_rank(name: str) -> int: """Unknown -> ValueError.""" # TODO raise NotImplementedError # ====================================================================================== # 4. Routing # ====================================================================================== @dataclass(frozen=True) class RoutingRule: """Match on what the caller IS, not on what model it wants. A caller that names a model has hard-coded a vendor decision into an agent.""" name: str deployments: Tuple[str, ...] # ordered: primary first, then fallbacks task_classes: Tuple[TaskClass, ...] = () tenants: Tuple[str, ...] = () classifications: Tuple[str, ...] = () # empty = any; else exact membership priority: int = 100 # lower wins def __post_init__(self) -> None: # TODO: validate every entry of `classifications` via classification_rank, so a # typo fails at construction rather than at 3 a.m. raise NotImplementedError def matches(self, request: NormalizedRequest) -> bool: """Empty tuple means "any". All three constraints must pass.""" # TODO raise NotImplementedError class Router: def __init__(self, deployments: Mapping[str, Deployment], rules: Sequence[RoutingRule]) -> None: """Reject (ValueError) a rule naming an unknown deployment. Store rules sorted by ``(priority, name)``.""" # TODO raise NotImplementedError def candidates(self, request: NormalizedRequest) -> List[Deployment]: """The ordered fallback chain. Take the FIRST matching rule whose chain is non-empty after applying the constraints a rule cannot express (see ``_admissible``). Return [] if none. """ # TODO raise NotImplementedError def _admissible(self, deployment: Deployment, request: NormalizedRequest) -> bool: """Residency (if not "any", the region must match) and the deployment's own classification ceiling.""" # TODO raise NotImplementedError # ====================================================================================== # 5. Rate limiting and quotas # ====================================================================================== class TokenBucket: def __init__(self, capacity: float, refill_per_second: float, *, now: Callable[[], float]) -> None: if capacity <= 0 or refill_per_second <= 0: raise ValueError("capacity and refill_per_second must be > 0") self.capacity = float(capacity) self.refill_per_second = float(refill_per_second) self.now = now self.tokens = float(capacity) self.last = now() def _refill(self) -> None: """``tokens = min(capacity, tokens + elapsed * rate)``; advance ``last``.""" # TODO raise NotImplementedError def try_consume(self, amount: float) -> bool: """Refill, then consume if possible. NEVER go negative. Negative amount -> ValueError.""" # TODO raise NotImplementedError def retry_after_seconds(self, amount: float) -> float: """``max(0, (amount - tokens) / rate)`` after refilling.""" # TODO raise NotImplementedError @dataclass class TenantLimits: """Both rate limits are needed: one request can be 100 000 tokens, so RPM alone does not protect the provider; TPM alone does not stop a flood of tiny requests.""" requests_per_minute: int tokens_per_minute: int monthly_budget_micros: int class RateLimiter: def __init__(self, limits: Mapping[str, TenantLimits], *, now: Callable[[], float]) -> None: self.limits = dict(limits) self.now = now self._rpm: Dict[str, TokenBucket] = {} self._tpm: Dict[str, TokenBucket] = {} def _buckets(self, tenant: str) -> Tuple[TokenBucket, TokenBucket]: """Lazily create both buckets (capacity = the per-minute limit, rate = limit/60). An unconfigured tenant -> QuotaExceeded.""" # TODO raise NotImplementedError def admit(self, tenant: str, estimated_tokens: int) -> None: """Raise RateLimited if either bucket refuses. Consume from BOTH only when both would admit — otherwise a request rejected on tokens still spends an RPM slot, and the tenant is throttled twice for one try. """ # TODO raise NotImplementedError class QuotaLedger: """Per-tenant monthly spend with a hard stop. Unbounded consumption is an availability risk as much as a budget one.""" def __init__(self, limits: Mapping[str, TenantLimits]) -> None: self.limits = dict(limits) self._spent: Dict[str, int] = {} def check(self, tenant: str) -> None: """Unknown tenant or spend >= budget -> QuotaExceeded (the boundary is inclusive).""" # TODO raise NotImplementedError def record(self, tenant: str, cost_micros: int) -> None: # TODO raise NotImplementedError def spent(self, tenant: str) -> int: return self._spent.get(tenant, 0) def remaining(self, tenant: str) -> int: """Never negative.""" # TODO raise NotImplementedError # ====================================================================================== # 6. Caching # ====================================================================================== def hash_embed(text: str, dimensions: int = 64) -> List[float]: """A deterministic hashing bag-of-words embedder, L2-normalized. For each lowercase whitespace token: blake2b digest, index = first 4 bytes mod dimensions, sign = +1 if digest[4] is even else -1. Normalize at the end; an all-zero vector stays all-zero (no division by zero). """ # TODO raise NotImplementedError def cosine(a: Sequence[float], b: Sequence[float]) -> float: """Dot product — both inputs are already normalized.""" # TODO raise NotImplementedError def cache_key(request: NormalizedRequest, deployment_name: str) -> str: """TENANT FIRST, then deployment, task class, temperature, max_output_tokens and the full prompt; hashed. Every cache key in a multi-tenant platform starts with the tenant. A cache that can cross a tenant boundary is a data breach with an excellent hit rate. """ # TODO raise NotImplementedError class ExactCache: def __init__(self, *, ttl_seconds: float, now: Callable[[], float], capacity: int = 1024) -> None: self.ttl = ttl_seconds self.now = now self.capacity = capacity self._entries: Dict[str, Tuple[float, NormalizedResponse]] = {} self.hits = 0 self.misses = 0 def get(self, key: str) -> Optional[NormalizedResponse]: """Expired entries are deleted and count as a miss. Update hits/misses.""" # TODO raise NotImplementedError def put(self, key: str, response: NormalizedResponse) -> None: """Evict the OLDEST entry when at capacity.""" # TODO raise NotImplementedError @dataclass class SemanticEntry: tenant: str vector: Tuple[float, ...] response: NormalizedResponse stored_at: float class SemanticCache: """Three non-negotiables: tenant-PARTITIONED (not merely filtered), a similarity floor, and never used for entitlement-dependent answers (the caller sets ``cacheable=False`` and the gateway obeys).""" def __init__(self, *, threshold: float, ttl_seconds: float, now: Callable[[], float], capacity_per_tenant: int = 256) -> None: if not 0.0 < threshold <= 1.0: raise ValueError("threshold must be in (0, 1]") self.threshold = threshold self.ttl = ttl_seconds self.now = now self.capacity_per_tenant = capacity_per_tenant self._by_tenant: Dict[str, List[SemanticEntry]] = {} self.hits = 0 self.misses = 0 def get(self, request: NormalizedRequest) -> Optional[NormalizedResponse]: """Drop expired entries, embed the prompt, take the BEST cosine score within this tenant's partition, and return it only if it meets the threshold.""" # TODO raise NotImplementedError def put(self, request: NormalizedRequest, response: NormalizedResponse) -> None: """Append to the tenant's partition, evicting the oldest at capacity.""" # TODO raise NotImplementedError # ====================================================================================== # 7. Accounting # ====================================================================================== @dataclass(frozen=True) class AccountingRecord: tenant: str agent_id: str deployment: str provider: str model: str task_class: TaskClass usage: Usage cost_micros: int latency_ms: int cache: str attempts: Tuple[str, ...] outcome: str # "ok" | the error class name class Accounting: """Records EVERY outcome, including failures — a cost model that counts only successes under-reports exactly during an incident.""" def __init__(self) -> None: self.records: List[AccountingRecord] = [] def record(self, record: AccountingRecord) -> None: self.records.append(record) def cost_by(self, key: str) -> Dict[str, int]: """Sum cost by one of tenant / agent_id / deployment / provider / model, returned sorted by key. Anything else -> ValueError.""" # TODO raise NotImplementedError def tokens_by_tenant(self) -> Dict[str, int]: # TODO raise NotImplementedError def cache_hit_rate(self) -> float: """Fraction of records whose cache is not "miss". No records -> 0.0.""" # TODO raise NotImplementedError def failover_rate(self) -> float: """Fraction of records with more than one attempt — the metric that shows a provider degrading before the error rate does.""" # TODO raise NotImplementedError # ====================================================================================== # 8. The gateway # ====================================================================================== def estimate_tokens(text: str) -> int: return (len(text) + 3) // 4 class Gateway: def __init__(self, *, router: Router, adapters: Mapping[str, ProviderAdapter], rate_limiter: RateLimiter, quotas: QuotaLedger, accounting: Accounting, now: Callable[[], float], exact_cache: Optional[ExactCache] = None, semantic_cache: Optional[SemanticCache] = None) -> None: self.router = router self.adapters = dict(adapters) self.rate_limiter = rate_limiter self.quotas = quotas self.accounting = accounting self.now = now self.exact_cache = exact_cache self.semantic_cache = semantic_cache def complete(self, request: NormalizedRequest) -> NormalizedResponse: """The full path, in the order that matters: 1. quota, then rate limit — cheapest checks first, and they protect capacity. Estimated tokens = ``estimate_tokens(prompt) + max_output_tokens``. On refusal: record the failure, then re-raise. 2. route. No candidates -> NoRouteAvailable (recorded). 3. cache lookup, keyed on the PRIMARY candidate. A hit is recorded with zero cost and returned. 4. execute with budget-aware fallback. 5. store in the cache. """ # TODO raise NotImplementedError def _elapsed_ms(self, started: float) -> int: # TODO raise NotImplementedError def _cache_lookup(self, request: NormalizedRequest, primary: Deployment) -> Optional[NormalizedResponse]: """Exact first, then semantic. Skip entirely when ``cacheable`` is False. A hit reports ``cost_micros=0`` and ``attempts=()``: cost means "what THIS call cost", and reporting the original double-counts spend. """ # TODO raise NotImplementedError def _cache_store(self, request: NormalizedRequest, primary: Deployment, response: NormalizedResponse) -> None: """Store in both caches — but NEVER when ``cacheable`` is False, and never when the finish reason is not STOP (a truncated or filtered answer is not the answer).""" # TODO raise NotImplementedError def _execute_with_fallback(self, request: NormalizedRequest, candidates: Sequence[Deployment], started: float) -> NormalizedResponse: """Walk the chain. Before EACH non-primary attempt, two independent refusals: - ``request.side_effecting`` -> re-raise the last error. A retried tool-executing call can double-execute. - remaining budget (``latency_budget_ms - elapsed``) < the candidate's ``expected_latency_ms`` -> BudgetExhausted. A fallback that does not fit the budget is decoration. On a provider error: if ``fall_over`` is False, record and re-raise immediately; otherwise remember it and continue. On success: compute cost from the deployment, record spend and accounting, and return with ``attempts`` and real latency filled in. If the chain exhausts, raise the last error. """ # TODO raise NotImplementedError def _record_failure(self, request: NormalizedRequest, exc: GatewayError, started: float, attempts: Tuple[str, ...]) -> None: """Append an AccountingRecord with zero cost and ``outcome=type(exc).__name__``.""" # TODO raise NotImplementedError # ====================================================================================== # 9. Provider adapters # ====================================================================================== def make_scripted_adapter( provider: str, *, reply: Callable[[NormalizedRequest], str], failures: Sequence[Optional[GatewayError]] = (), input_tokens: Optional[Callable[[NormalizedRequest], int]] = None, output_tokens: int = 64, cached_fraction: float = 0.0, finish_reason: FinishReason = FinishReason.STOP, ) -> ProviderAdapter: """A deterministic adapter. ``failures`` is consumed one entry per call: a GatewayError is raised, None succeeds. That is how a test scripts "the primary 429s once, then works". """ # TODO raise NotImplementedError def main() -> None: print("implement the TODOs, then compare with `python solution.py`") if __name__ == "__main__": main()