diff --git a/README.md b/README.md index c113725..e9751b2 100644 --- a/README.md +++ b/README.md @@ -80,6 +80,19 @@ in-cluster scrapers can read it; restrict it with network policy rather than a b | `router_cache_events_total` | Cache lookups by cache and result. | | `router_model_load_seconds` | Model load and cold-start duration. | +## Caching + +Cache keys bind the tenant, the active model-catalog fingerprint, and the generation +parameters, so promoting any model revision or routing policy version invalidates every +dependent entry. Responses carry `X-Cache: miss | exact | semantic`. + +| Cache | Eligibility | +|---|---| +| Exact response | Deterministic requests (`temperature = 0`) that are not `restricted`. | +| Prefix | Reported per model revision for repeated instruction prefixes. | +| Semantic | Disabled by default; requires `public` privacy, deterministic generation, and an extraction, classification, or summarization task. | +| Router decision | Reuses stable task classification; cleared when the policy version changes. | + ## Runtime settings All settings use the `ROUTER_` prefix. @@ -91,6 +104,12 @@ All settings use the `ROUTER_` prefix. | `ROUTER_ADMISSION_TIMEOUT_SECONDS` | `0.25` | Time allowed to wait for capacity. | | `ROUTER_QUOTA_REQUESTS_PER_MINUTE` | `120` | Per-token sliding-window quota. | | `ROUTER_EXTERNAL_FALLBACK_ENABLED` | `false` | Operator gate for external fallback. | +| `ROUTER_ROUTING_POLICY_VERSION` | `v1` | Invalidates router and response caches when changed. | +| `ROUTER_CACHE_ENABLED` | `true` | Master switch for all cache tiers. | +| `ROUTER_CACHE_TTL_SECONDS` | `300` | Exact-response entry lifetime. | +| `ROUTER_CACHE_MAX_ENTRIES` | `1024` | Bound on cached responses. | +| `ROUTER_SEMANTIC_CACHE_ENABLED` | `false` | Enables similarity reuse for approved tasks. | +| `ROUTER_SEMANTIC_SIMILARITY_THRESHOLD` | `0.92` | Minimum similarity for a semantic hit. | External routing also requires public data and request-level opt-in. Private and restricted requests are never eligible for an external route. diff --git a/src/llm_router/app.py b/src/llm_router/app.py index 1ce0b33..fd257b6 100644 --- a/src/llm_router/app.py +++ b/src/llm_router/app.py @@ -14,13 +14,26 @@ QuotaExceededError, SlidingWindowQuota, ) -from llm_router.backends import InferenceBackend, MockInferenceBackend +from llm_router.backends import BackendResult, InferenceBackend, MockInferenceBackend +from llm_router.caching import ( + CachedCompletion, + CacheStore, + InMemoryCacheStore, + PrefixTracker, + RouterDecisionCache, + SemanticCache, + build_cache_key, + catalog_fingerprint, + exact_cache_eligible, + semantic_cache_eligible, +) from llm_router.config import Settings, get_settings from llm_router.models import ( ChatCompletionChoice, ChatCompletionRequest, ChatCompletionResponse, ChatMessage, + RouteDecision, Usage, ) from llm_router.observability import Metrics @@ -32,6 +45,7 @@ def create_app( *, backend: InferenceBackend | None = None, metrics: Metrics | None = None, + cache_store: CacheStore | None = None, ) -> FastAPI: runtime_settings = settings or get_settings() router = Router( @@ -45,6 +59,17 @@ def create_app( quota = SlidingWindowQuota(runtime_settings.quota_requests_per_minute) inference_backend = backend or MockInferenceBackend() telemetry = metrics if metrics is not None else Metrics() + exact_cache: CacheStore = cache_store or InMemoryCacheStore( + max_entries=runtime_settings.cache_max_entries, + ttl_seconds=runtime_settings.cache_ttl_seconds, + ) + semantic_cache = SemanticCache(threshold=runtime_settings.semantic_similarity_threshold) + decision_cache = RouterDecisionCache(policy_version=runtime_settings.routing_policy_version) + prefix_tracker = PrefixTracker() + catalog_version = catalog_fingerprint( + (profile.revision for profile in router.profiles), + runtime_settings.routing_policy_version, + ) @asynccontextmanager async def lifespan(app: FastAPI) -> AsyncIterator[None]: @@ -131,6 +156,80 @@ async def models() -> dict[str, object]: ] return {"object": "list", "data": visible} + def _completion_response( + *, + model_id: str, + text: str, + prompt_tokens: int, + completion_tokens: int, + routing: dict[str, object], + finish_reason: str = "stop", + ) -> ChatCompletionResponse: + return ChatCompletionResponse( + id=f"chatcmpl-{uuid.uuid4().hex}", + created=int(time.time()), + model=model_id, + choices=[ + ChatCompletionChoice( + message=ChatMessage(role="assistant", content=text), + finish_reason="length" if finish_reason == "length" else "stop", + ) + ], + usage=Usage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=prompt_tokens + completion_tokens, + ), + routing=routing, + ) + + async def _lookup_cache( + payload: ChatCompletionRequest, prompt: str, cache_key: str, tenant: str + ) -> tuple[str, CachedCompletion] | None: + if not runtime_settings.cache_enabled: + return None + if exact_cache_eligible(payload): + entry = await exact_cache.get(cache_key) + telemetry.record_cache_event("exact", "hit" if entry is not None else "miss") + if entry is not None: + return "exact", entry + if ( + runtime_settings.semantic_cache_enabled + and payload.routing.task is not None + and semantic_cache_eligible(payload, payload.routing.task) + ): + scope = semantic_cache.scope(payload, tenant=tenant, model_revision=catalog_version) + match = semantic_cache.lookup(scope, prompt) + telemetry.record_cache_event("semantic", "hit" if match is not None else "miss") + if match is not None: + return "semantic", match + return None + + async def _store_cache( + payload: ChatCompletionRequest, + prompt: str, + cache_key: str, + tenant: str, + decision: RouteDecision, + result: BackendResult, + ) -> None: + if not runtime_settings.cache_enabled: + return + entry = CachedCompletion( + text=result.text, + model_id=decision.profile.id, + model_revision=decision.profile.revision, + prompt_tokens=result.prompt_tokens, + completion_tokens=result.completion_tokens, + ) + if exact_cache_eligible(payload): + await exact_cache.set(cache_key, entry) + if runtime_settings.semantic_cache_enabled and semantic_cache_eligible( + payload, decision.task + ): + scope = semantic_cache.scope(payload, tenant=tenant, model_revision=catalog_version) + semantic_cache.store(scope, prompt, entry) + @app.post("/v1/chat/completions", response_model=ChatCompletionResponse) async def chat_completions( payload: ChatCompletionRequest, @@ -139,8 +238,41 @@ async def chat_completions( ) -> ChatCompletionResponse: started = time.perf_counter() await quota.consume(subject) - decision = router.select(payload) + prompt = payload.prompt + cache_key = build_cache_key( + payload, tenant=subject, model_revision=catalog_version, prompt=prompt + ) + + cached = await _lookup_cache(payload, prompt, cache_key, subject) + if cached is not None: + hit_name, entry = cached + response.headers["X-Cache"] = hit_name + response.headers["X-Route-Model"] = entry.model_id + response.headers["X-Route-Revision"] = entry.model_revision + return _completion_response( + model_id=entry.model_id, + text=entry.text, + prompt_tokens=entry.prompt_tokens, + completion_tokens=entry.completion_tokens, + routing={ + "model_revision": entry.model_revision, + "cache": hit_name, + "reason": f"served from the {hit_name} cache", + }, + ) + + 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) + if runtime_settings.cache_enabled: + decision_cache.set(prompt, decision.task) telemetry.record_route(decision, privacy=payload.routing.privacy.value) + telemetry.record_cache_event( + "prefix", + "hit" + if prefix_tracker.observe(prompt, model_revision=decision.profile.revision) + else "miss", + ) telemetry.queued_requests.inc() try: @@ -164,6 +296,9 @@ async def chat_completions( completion_tokens=result.completion_tokens, ) + await _store_cache(payload, prompt, cache_key, subject, decision, result) + + response.headers["X-Cache"] = "miss" response.headers["X-Route-Model"] = decision.profile.id response.headers["X-Route-Revision"] = decision.profile.revision response.headers["X-Route-Reason"] = decision.reason diff --git a/src/llm_router/caching.py b/src/llm_router/caching.py new file mode 100644 index 0000000..f4c6276 --- /dev/null +++ b/src/llm_router/caching.py @@ -0,0 +1,216 @@ +"""Cache tiers defined in section 12 of the design specification. + +Every cache key includes the tenant, the immutable model revision, and the +generation parameters, so a cache entry can never cross a tenant boundary or +survive a model promotion. Semantic caching is additionally restricted to +non-sensitive, deterministic task classes. +""" + +import hashlib +import time +from collections import OrderedDict +from collections.abc import Iterable +from dataclasses import dataclass, field +from typing import Protocol + +from llm_router.models import ChatCompletionRequest, PrivacyClass, TaskClass + +SEMANTIC_CACHE_ELIGIBLE_TASKS = frozenset( + {TaskClass.CLASSIFICATION, TaskClass.EXTRACTION, TaskClass.SUMMARIZATION} +) + + +@dataclass(frozen=True) +class CachedCompletion: + text: str + model_id: str + model_revision: str + prompt_tokens: int + completion_tokens: int + + +class CacheStore(Protocol): + """Minimal async key-value contract implemented by memory and Redis stores.""" + + async def get(self, key: str) -> CachedCompletion | None: ... + + async def set(self, key: str, value: CachedCompletion) -> None: ... + + +class InMemoryCacheStore: + """Bounded TTL store used for single-replica deployments and tests.""" + + def __init__(self, *, max_entries: int = 1024, ttl_seconds: float = 300.0) -> None: + self._max_entries = max_entries + self._ttl_seconds = ttl_seconds + self._entries: OrderedDict[str, tuple[float, CachedCompletion]] = OrderedDict() + + async def get(self, key: str) -> CachedCompletion | None: + entry = self._entries.get(key) + if entry is None: + return None + expires_at, value = entry + if expires_at <= time.monotonic(): + del self._entries[key] + return None + self._entries.move_to_end(key) + return value + + async def set(self, key: str, value: CachedCompletion) -> None: + self._entries[key] = (time.monotonic() + self._ttl_seconds, value) + self._entries.move_to_end(key) + while len(self._entries) > self._max_entries: + self._entries.popitem(last=False) + + def __len__(self) -> int: + return len(self._entries) + + +def _tokenize(text: str) -> frozenset[str]: + return frozenset(word for word in text.lower().split() if word) + + +def similarity(left: str, right: str) -> float: + """Jaccard similarity over normalized tokens, used for semantic lookups.""" + + left_tokens, right_tokens = _tokenize(left), _tokenize(right) + if not left_tokens or not right_tokens: + return 0.0 + return len(left_tokens & right_tokens) / len(left_tokens | right_tokens) + + +def build_cache_key( + request: ChatCompletionRequest, + *, + tenant: str, + model_revision: str, + prompt: str | None = None, +) -> str: + """Derive an exact-response key from tenant, revision, parameters, and prompt.""" + + material = "|".join( + ( + tenant, + model_revision, + f"max_tokens={request.max_tokens}", + f"temperature={request.temperature}", + f"privacy={request.routing.privacy.value}", + prompt if prompt is not None else request.prompt, + ) + ) + return hashlib.sha256(material.encode()).hexdigest() + + +def prefix_key(prompt: str, *, model_revision: str, prefix_chars: int = 512) -> str: + """Key for repeated instruction and context prefixes bound to one revision.""" + + return hashlib.sha256(f"{model_revision}|{prompt[:prefix_chars]}".encode()).hexdigest() + + +def exact_cache_eligible(request: ChatCompletionRequest) -> bool: + """Exact reuse requires deterministic generation and a non-restricted class.""" + + return request.temperature == 0.0 and request.routing.privacy != PrivacyClass.RESTRICTED + + +def semantic_cache_eligible(request: ChatCompletionRequest, task: TaskClass) -> bool: + """Semantic reuse is limited to public, deterministic, approved task classes.""" + + return ( + request.temperature == 0.0 + and request.routing.privacy == PrivacyClass.PUBLIC + and task in SEMANTIC_CACHE_ELIGIBLE_TASKS + ) + + +@dataclass +class SemanticCache: + """Similarity cache scoped by tenant, revision, and generation parameters.""" + + threshold: float = 0.92 + max_entries: int = 256 + _entries: OrderedDict[str, list[tuple[str, CachedCompletion]]] = field( + default_factory=OrderedDict + ) + + def scope(self, request: ChatCompletionRequest, *, tenant: str, model_revision: str) -> str: + return build_cache_key(request, tenant=tenant, model_revision=model_revision, prompt="") + + def lookup(self, scope: str, prompt: str) -> CachedCompletion | None: + candidates: Iterable[tuple[str, CachedCompletion]] = self._entries.get(scope, []) + best: tuple[float, CachedCompletion] | None = None + for stored_prompt, value in candidates: + score = similarity(stored_prompt, prompt) + if score >= self.threshold and (best is None or score > best[0]): + best = (score, value) + return None if best is None else best[1] + + def store(self, scope: str, prompt: str, value: CachedCompletion) -> None: + bucket = self._entries.setdefault(scope, []) + bucket.append((prompt, value)) + self._entries.move_to_end(scope) + while len(self._entries) > self.max_entries: + self._entries.popitem(last=False) + + +@dataclass +class RouterDecisionCache: + """Caches stable task classifications and invalidates on policy changes.""" + + policy_version: str + max_entries: int = 512 + _entries: OrderedDict[str, TaskClass] = field(default_factory=OrderedDict) + + def _key(self, prompt: str) -> str: + return hashlib.sha256(f"{self.policy_version}|{prompt}".encode()).hexdigest() + + def get(self, prompt: str) -> TaskClass | None: + key = self._key(prompt) + value = self._entries.get(key) + if value is not None: + self._entries.move_to_end(key) + return value + + def set(self, prompt: str, task: TaskClass) -> None: + key = self._key(prompt) + self._entries[key] = task + self._entries.move_to_end(key) + while len(self._entries) > self.max_entries: + self._entries.popitem(last=False) + + def invalidate(self, policy_version: str) -> None: + if policy_version != self.policy_version: + self.policy_version = policy_version + self._entries.clear() + + def __len__(self) -> int: + return len(self._entries) + + +def catalog_fingerprint(revisions: Iterable[str], policy_version: str) -> str: + """Stable identity for the active model catalog and routing policy. + + Cache lookups happen before routing, so keys are bound to this fingerprint: + promoting any model or policy version invalidates every dependent entry. + """ + + material = "|".join((policy_version, *sorted(revisions))) + return hashlib.sha256(material.encode()).hexdigest()[:16] + + +class PrefixTracker: + """Tracks reuse of repeated instruction prefixes for cache-hit reporting.""" + + def __init__(self, *, max_entries: int = 1024, prefix_chars: int = 512) -> None: + self._max_entries = max_entries + self._prefix_chars = prefix_chars + self._seen: OrderedDict[str, None] = OrderedDict() + + def observe(self, prompt: str, *, model_revision: str) -> bool: + key = prefix_key(prompt, model_revision=model_revision, prefix_chars=self._prefix_chars) + hit = key in self._seen + self._seen[key] = None + self._seen.move_to_end(key) + while len(self._seen) > self._max_entries: + self._seen.popitem(last=False) + return hit diff --git a/src/llm_router/config.py b/src/llm_router/config.py index a0822fb..163c124 100644 --- a/src/llm_router/config.py +++ b/src/llm_router/config.py @@ -15,6 +15,12 @@ class Settings(BaseSettings): admission_timeout_seconds: float = Field(default=0.25, gt=0) quota_requests_per_minute: int = Field(default=120, ge=1) external_fallback_enabled: bool = False + routing_policy_version: str = "v1" + cache_enabled: bool = True + cache_ttl_seconds: float = Field(default=300.0, gt=0) + cache_max_entries: int = Field(default=1024, ge=1) + semantic_cache_enabled: bool = False + semantic_similarity_threshold: float = Field(default=0.92, ge=0.0, le=1.0) @model_validator(mode="after") def reject_development_key_in_shared_environments(self) -> "Settings": diff --git a/src/llm_router/routing.py b/src/llm_router/routing.py index ad3c409..9d21dcb 100644 --- a/src/llm_router/routing.py +++ b/src/llm_router/routing.py @@ -89,8 +89,10 @@ def classify_task(self, request: ChatCompletionRequest) -> TaskClass: TaskClass.GENERAL, ) - def select(self, request: ChatCompletionRequest) -> RouteDecision: - task = self.classify_task(request) + def select( + self, request: ChatCompletionRequest, *, task: TaskClass | 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 candidates = [ diff --git a/tests/e2e/api.spec.ts b/tests/e2e/api.spec.ts index 6293f7e..8cbe04a 100644 --- a/tests/e2e/api.spec.ts +++ b/tests/e2e/api.spec.ts @@ -62,3 +62,35 @@ test("publishes scrape-ready metrics for a completed request", async ({ request expect(body).toContain("router_tokens_total"); expect(body).toContain("router_predicted_quality_sum"); }); + +test("serves a repeated deterministic request from the exact cache", async ({ request }) => { + const body = { + model: "auto", + messages: [{ role: "user", content: "Classify this end-to-end cache probe" }], + routing: { privacy: "public" }, + }; + + const first = await request.post("/v1/chat/completions", { data: body }); + expect(first.status()).toBe(200); + expect(first.headers()["x-cache"]).toBe("miss"); + + const second = await request.post("/v1/chat/completions", { data: body }); + expect(second.status()).toBe(200); + expect(second.headers()["x-cache"]).toBe("exact"); + expect((await second.json()).routing.cache).toBe("exact"); + expect(second.headers()["x-route-model"]).toBe(first.headers()["x-route-model"]); +}); + +test("never caches restricted-class requests", async ({ request }) => { + const body = { + model: "auto", + messages: [{ role: "user", content: "Extract fields from this restricted record" }], + routing: { privacy: "restricted" }, + }; + + await request.post("/v1/chat/completions", { data: body }); + const repeat = await request.post("/v1/chat/completions", { data: body }); + + expect(repeat.status()).toBe(200); + expect(repeat.headers()["x-cache"]).toBe("miss"); +}); diff --git a/tests/integration/test_caching_api.py b/tests/integration/test_caching_api.py new file mode 100644 index 0000000..10f496a --- /dev/null +++ b/tests/integration/test_caching_api.py @@ -0,0 +1,127 @@ +from fastapi.testclient import TestClient + +from llm_router.app import create_app +from llm_router.config import Settings + +HEADERS = {"Authorization": "Bearer cache-key"} + + +def build_client(**overrides: object) -> TestClient: + settings = Settings(api_keys="cache-key", **overrides) # type: ignore[arg-type] + return TestClient(create_app(settings)) + + +def completion(client: TestClient, **body: object) -> tuple[int, dict[str, str], dict[str, object]]: + payload: dict[str, object] = { + "model": "auto", + "messages": [{"role": "user", "content": "Classify this refund request"}], + "routing": {"privacy": "public"}, + } + payload.update(body) + response = client.post("/v1/chat/completions", headers=HEADERS, json=payload) + return response.status_code, dict(response.headers), response.json() + + +def test_repeated_deterministic_request_is_served_from_the_exact_cache() -> None: + with build_client() as client: + first_status, first_headers, first_body = completion(client) + second_status, second_headers, second_body = completion(client) + + assert (first_status, second_status) == (200, 200) + assert first_headers["x-cache"] == "miss" + assert second_headers["x-cache"] == "exact" + assert second_headers["x-route-model"] == first_headers["x-route-model"] + assert second_body["routing"]["cache"] == "exact" + assert ( + second_body["choices"][0]["message"]["content"] + == first_body["choices"][0]["message"]["content"] + ) + + +def test_restricted_requests_are_never_served_from_cache() -> None: + with build_client() as client: + completion(client, routing={"privacy": "restricted"}) + _, headers, body = completion(client, routing={"privacy": "restricted"}) + + assert headers["x-cache"] == "miss" + assert "cache" not in body["routing"] + + +def test_sampled_requests_are_not_cached() -> None: + with build_client() as client: + completion(client, temperature=0.7) + _, headers, _ = completion(client, temperature=0.7) + + assert headers["x-cache"] == "miss" + + +def test_cache_can_be_disabled_by_configuration() -> None: + with build_client(cache_enabled=False) as client: + completion(client) + _, headers, _ = completion(client) + metrics = client.get("/metrics").text + + assert headers["x-cache"] == "miss" + assert 'router_cache_events_total{cache="exact"' not in metrics + + +def test_semantic_cache_serves_similar_public_prompts_when_enabled() -> None: + with build_client(semantic_cache_enabled=True, semantic_similarity_threshold=0.5) as client: + client.post( + "/v1/chat/completions", + headers=HEADERS, + json={ + "model": "auto", + "messages": [{"role": "user", "content": "classify this refund request please"}], + "routing": {"privacy": "public", "task": "classification"}, + }, + ) + response = client.post( + "/v1/chat/completions", + headers=HEADERS, + json={ + "model": "auto", + "messages": [{"role": "user", "content": "classify this refund request now"}], + "routing": {"privacy": "public", "task": "classification"}, + }, + ) + + assert response.status_code == 200 + assert response.headers["x-cache"] == "semantic" + assert response.json()["routing"]["cache"] == "semantic" + + +def test_semantic_cache_is_not_used_for_private_requests() -> None: + with build_client(semantic_cache_enabled=True, semantic_similarity_threshold=0.1) as client: + client.post( + "/v1/chat/completions", + headers=HEADERS, + json={ + "model": "auto", + "messages": [{"role": "user", "content": "classify this private ticket"}], + "routing": {"privacy": "private", "task": "classification"}, + }, + ) + response = client.post( + "/v1/chat/completions", + headers=HEADERS, + json={ + "model": "auto", + "messages": [{"role": "user", "content": "classify this private matter"}], + "routing": {"privacy": "private", "task": "classification"}, + }, + ) + + assert response.headers["x-cache"] == "miss" + + +def test_cache_metrics_report_router_prefix_and_exact_outcomes() -> None: + with build_client() as client: + completion(client) + completion(client) + metrics = client.get("/metrics").text + + assert 'router_cache_events_total{cache="exact",result="hit"} 1.0' in metrics + assert 'router_cache_events_total{cache="exact",result="miss"} 1.0' in metrics + assert 'router_cache_events_total{cache="router",result="miss"} 1.0' in metrics + assert 'router_cache_events_total{cache="prefix",result="miss"} 1.0' in metrics diff --git a/tests/unit/test_caching.py b/tests/unit/test_caching.py new file mode 100644 index 0000000..7003419 --- /dev/null +++ b/tests/unit/test_caching.py @@ -0,0 +1,151 @@ +import pytest + +from llm_router.caching import ( + CachedCompletion, + InMemoryCacheStore, + RouterDecisionCache, + SemanticCache, + build_cache_key, + exact_cache_eligible, + prefix_key, + semantic_cache_eligible, + similarity, +) +from llm_router.models import ChatCompletionRequest, ChatMessage, PrivacyClass, TaskClass + + +def make_request(**routing: object) -> ChatCompletionRequest: + return ChatCompletionRequest( + messages=[ChatMessage(role="user", content="classify this support ticket")], + temperature=float(routing.pop("temperature", 0.0)), + routing=dict(routing), # type: ignore[arg-type] + ) + + +def sample_value() -> CachedCompletion: + return CachedCompletion( + text="cached", + model_id="small-specialist", + model_revision="mock-small@sha256:dev", + prompt_tokens=4, + completion_tokens=2, + ) + + +@pytest.mark.asyncio +async def test_memory_store_returns_stored_value_then_expires() -> None: + store = InMemoryCacheStore(ttl_seconds=-1.0) + await store.set("key", sample_value()) + + assert await store.get("key") is None + assert len(store) == 0 + + +@pytest.mark.asyncio +async def test_memory_store_evicts_least_recently_used_entries() -> None: + store = InMemoryCacheStore(max_entries=2) + await store.set("a", sample_value()) + await store.set("b", sample_value()) + await store.get("a") + await store.set("c", sample_value()) + + assert await store.get("b") is None + assert await store.get("a") is not None + assert await store.get("c") is not None + + +@pytest.mark.asyncio +async def test_memory_store_misses_unknown_key() -> None: + store = InMemoryCacheStore() + + assert await store.get("absent") is None + + +def test_cache_key_separates_tenants_revisions_and_parameters() -> None: + request = make_request() + base = build_cache_key(request, tenant="tenant-a", model_revision="rev-1") + + assert base != build_cache_key(request, tenant="tenant-b", model_revision="rev-1") + assert base != build_cache_key(request, tenant="tenant-a", model_revision="rev-2") + assert base == build_cache_key(request, tenant="tenant-a", model_revision="rev-1") + + +def test_cache_key_changes_with_generation_parameters() -> None: + request = make_request() + hotter = ChatCompletionRequest(messages=request.messages, temperature=0.7) + + assert build_cache_key(request, tenant="t", model_revision="r") != build_cache_key( + hotter, tenant="t", model_revision="r" + ) + + +def test_prefix_key_is_revision_bound_and_prefix_scoped() -> None: + shared = "system instructions " * 40 + assert prefix_key(shared + "tail-a", model_revision="rev-1", prefix_chars=32) == prefix_key( + shared + "tail-b", model_revision="rev-1", prefix_chars=32 + ) + assert prefix_key(shared, model_revision="rev-1") != prefix_key(shared, model_revision="rev-2") + + +def test_exact_cache_rejects_sampling_and_restricted_data() -> None: + assert exact_cache_eligible(make_request()) + assert not exact_cache_eligible(make_request(temperature=0.5)) + assert not exact_cache_eligible(make_request(privacy=PrivacyClass.RESTRICTED)) + + +def test_semantic_cache_requires_public_deterministic_approved_task() -> None: + public = make_request(privacy=PrivacyClass.PUBLIC) + + assert semantic_cache_eligible(public, TaskClass.CLASSIFICATION) + assert not semantic_cache_eligible(public, TaskClass.REASONING) + assert not semantic_cache_eligible(make_request(), TaskClass.CLASSIFICATION) + assert not semantic_cache_eligible( + make_request(privacy=PrivacyClass.PUBLIC, temperature=0.9), TaskClass.CLASSIFICATION + ) + + +def test_similarity_scores_identical_and_disjoint_prompts() -> None: + assert similarity("route this ticket", "route this ticket") == 1.0 + assert similarity("route this ticket", "unrelated words entirely") == 0.0 + assert similarity("", "anything") == 0.0 + + +def test_semantic_cache_returns_best_match_above_threshold_only() -> None: + cache = SemanticCache(threshold=0.6) + cache.store("scope", "classify this support ticket", sample_value()) + + assert cache.lookup("scope", "classify this support ticket now") is not None + assert cache.lookup("scope", "translate this document") is None + assert cache.lookup("other-scope", "classify this support ticket") is None + + +def test_semantic_cache_scope_isolates_tenants() -> None: + cache = SemanticCache() + request = make_request(privacy=PrivacyClass.PUBLIC) + + assert cache.scope(request, tenant="a", model_revision="r") != cache.scope( + request, tenant="b", model_revision="r" + ) + + +def test_router_cache_reuses_classification_until_policy_changes() -> None: + cache = RouterDecisionCache(policy_version="v1") + cache.set("classify this", TaskClass.CLASSIFICATION) + + assert cache.get("classify this") is TaskClass.CLASSIFICATION + + cache.invalidate("v1") + assert cache.get("classify this") is TaskClass.CLASSIFICATION + + cache.invalidate("v2") + assert cache.get("classify this") is None + assert len(cache) == 0 + + +def test_router_cache_bounds_entry_count() -> None: + cache = RouterDecisionCache(policy_version="v1", max_entries=2) + for index in range(5): + cache.set(f"prompt-{index}", TaskClass.GENERAL) + + assert len(cache) == 2 + assert cache.get("prompt-0") is None