"""Registry resolver for the inference gateway. Resolves agent stage bindings to complete InferenceTarget instances using TTL-cached lookups against the registry tables. Auth secrets are NOT resolved during caching — only at invocation time. Requirements: 3.5, 3.9 """ from __future__ import annotations import logging import time from typing import Any, TypeVar from uuid import UUID from services.shared.inference.errors import InferenceError, InferenceErrorCategory from services.shared.inference.models import InferenceTarget, ProviderCapabilities logger = logging.getLogger(__name__) T = TypeVar("T") # Default cache TTL in seconds DEFAULT_CACHE_TTL_SECONDS = 60.0 # --------------------------------------------------------------------------- # RegistryCache — internal TTL dict # --------------------------------------------------------------------------- class RegistryCache: """TTL-based in-memory cache for registry lookups. Keys follow the format: - "binding:{agent_id}:{stage}" - "endpoint:{endpoint_id}" - "deployment:{deployment_id}" Expired entries are lazily evicted on access. """ def __init__(self, ttl_seconds: float = DEFAULT_CACHE_TTL_SECONDS) -> None: self._ttl_seconds = ttl_seconds self._store: dict[str, tuple[float, Any]] = {} @property def ttl_seconds(self) -> float: return self._ttl_seconds def get(self, key: str) -> Any | None: """Retrieve a cached value if it exists and hasn't expired.""" entry = self._store.get(key) if entry is None: return None stored_at, value = entry if (time.monotonic() - stored_at) > self._ttl_seconds: del self._store[key] return None return value def set(self, key: str, value: Any) -> None: """Store a value with the current timestamp.""" self._store[key] = (time.monotonic(), value) def invalidate(self, key_pattern: str) -> None: """Invalidate all entries whose key starts with the given pattern. Supports prefix-based invalidation: - invalidate("endpoint:abc-123") removes that specific endpoint - invalidate("endpoint:") removes ALL endpoint entries - invalidate("binding:agent-1:") removes all bindings for agent-1 """ keys_to_remove = [k for k in self._store if k.startswith(key_pattern)] for k in keys_to_remove: del self._store[k] def clear(self) -> None: """Clear all cached entries.""" self._store.clear() def __len__(self) -> int: return len(self._store) def __contains__(self, key: str) -> bool: """Check if a non-expired entry exists for the key.""" return self.get(key) is not None # --------------------------------------------------------------------------- # Database query protocol # --------------------------------------------------------------------------- class RegistryDB: """Protocol for registry database queries. In production this would be backed by asyncpg. For testing, a simple dict-based mock implements the same interface. """ async def get_active_binding( self, agent_id: UUID, stage: str ) -> dict[str, Any] | None: """Get the active binding for an agent+stage. Returns a dict with keys: id, agent_id, stage, model_deployment_id, route_order, routing_config, is_active, revision. Returns None if no active binding exists. """ raise NotImplementedError async def get_model_deployment( self, deployment_id: UUID ) -> dict[str, Any] | None: """Get a model deployment by ID. Returns a dict with keys: id, endpoint_id, served_model_name, display_name, capabilities, context_window, max_output_tokens, quantization, runtime_metadata, enabled, revision. """ raise NotImplementedError async def get_inference_endpoint( self, endpoint_id: UUID ) -> dict[str, Any] | None: """Get an inference endpoint by ID. Returns a dict with keys: id, name, protocol, base_url, auth_secret_ref, auth_scheme, default_headers, health_path, enabled, revision. """ raise NotImplementedError # --------------------------------------------------------------------------- # RegistryResolver # --------------------------------------------------------------------------- class RegistryResolver: """Resolves agent stage bindings to InferenceTarget instances. Resolution order: 1. Find active binding for (agent_id, stage) 2. Get model_deployment from binding 3. Get inference_endpoint from deployment 4. Build InferenceTarget Fail-closed: any missing or disabled resource raises InferenceError with CAPABILITY_UNAVAILABLE. Never returns a fallback target. Auth secrets are NOT resolved during caching — auth_secret_ref is preserved as-is in the target. Resolution happens at invocation time by the client layer. """ def __init__( self, db: RegistryDB, *, cache_ttl_seconds: float = DEFAULT_CACHE_TTL_SECONDS, ) -> None: self._db = db self._cache = RegistryCache(ttl_seconds=cache_ttl_seconds) @property def cache(self) -> RegistryCache: """Access the internal cache (for testing/inspection).""" return self._cache async def resolve_target(self, agent_id: UUID, stage: str) -> InferenceTarget: """Resolve the active binding for an agent+stage into an InferenceTarget. Uses TTL-cached lookups. If any component is missing or disabled, raises InferenceError(CAPABILITY_UNAVAILABLE). Auth secret is NOT resolved here — only at invocation time. """ # Check cache first for the full resolved target cache_key = f"binding:{agent_id}:{stage}" cached_target = self._cache.get(cache_key) if cached_target is not None: return cached_target # Step 1: Find active binding binding = await self._db.get_active_binding(agent_id, stage) if binding is None: raise InferenceError( InferenceErrorCategory.CAPABILITY_UNAVAILABLE, f"No active binding for agent={agent_id} stage={stage}", ) if not binding.get("is_active", False): raise InferenceError( InferenceErrorCategory.CAPABILITY_UNAVAILABLE, f"Binding for agent={agent_id} stage={stage} is inactive", ) deployment_id = binding.get("model_deployment_id") if deployment_id is None: raise InferenceError( InferenceErrorCategory.CAPABILITY_UNAVAILABLE, f"Binding for agent={agent_id} stage={stage} has no deployment", ) # Step 2: Get model deployment deployment = await self._resolve_deployment(deployment_id) # Step 3: Get inference endpoint endpoint_id = deployment["endpoint_id"] endpoint = await self._resolve_endpoint(endpoint_id) # Step 4: Build InferenceTarget capabilities_data = deployment.get("capabilities", {}) capabilities = ProviderCapabilities( chat_completions=capabilities_data.get("chat_completions", False), responses_api=capabilities_data.get("responses_api", False), json_schema=capabilities_data.get("json_schema", False), json_object=capabilities_data.get("json_object", False), seed=capabilities_data.get("seed", False), usage=capabilities_data.get("usage", False), max_completion_tokens=capabilities_data.get("max_completion_tokens", False), reasoning_toggle=capabilities_data.get("reasoning_toggle", False), model_listing=capabilities_data.get("model_listing", False), ) extra_headers = endpoint.get("default_headers", {}) if not isinstance(extra_headers, dict): extra_headers = {} runtime_metadata = deployment.get("runtime_metadata", {}) extra_body = runtime_metadata.get("extra_body", {}) if isinstance(runtime_metadata, dict) else {} target = InferenceTarget( endpoint_id=endpoint_id, deployment_id=deployment_id, protocol=endpoint["protocol"], base_url=endpoint["base_url"], model=deployment["served_model_name"], capabilities=capabilities, auth_secret_ref=endpoint.get("auth_secret_ref"), auth_scheme=endpoint.get("auth_scheme", "bearer"), extra_headers=extra_headers, extra_body=extra_body, context_window=deployment.get("context_window") or 0, max_output_tokens=deployment.get("max_output_tokens"), ) # Cache the resolved target self._cache.set(cache_key, target) return target async def _resolve_deployment(self, deployment_id: UUID) -> dict[str, Any]: """Resolve a model deployment, using cache if available.""" dep_cache_key = f"deployment:{deployment_id}" cached = self._cache.get(dep_cache_key) if cached is not None: return cached deployment = await self._db.get_model_deployment(deployment_id) if deployment is None: raise InferenceError( InferenceErrorCategory.CAPABILITY_UNAVAILABLE, f"Model deployment {deployment_id} not found", ) if not deployment.get("enabled", False): raise InferenceError( InferenceErrorCategory.CAPABILITY_UNAVAILABLE, f"Model deployment {deployment_id} is disabled", ) self._cache.set(dep_cache_key, deployment) return deployment async def _resolve_endpoint(self, endpoint_id: UUID) -> dict[str, Any]: """Resolve an inference endpoint, using cache if available.""" ep_cache_key = f"endpoint:{endpoint_id}" cached = self._cache.get(ep_cache_key) if cached is not None: return cached endpoint = await self._db.get_inference_endpoint(endpoint_id) if endpoint is None: raise InferenceError( InferenceErrorCategory.CAPABILITY_UNAVAILABLE, f"Inference endpoint {endpoint_id} not found", ) if not endpoint.get("enabled", False): raise InferenceError( InferenceErrorCategory.CAPABILITY_UNAVAILABLE, f"Inference endpoint {endpoint_id} is disabled", ) self._cache.set(ep_cache_key, endpoint) return endpoint def invalidate(self, endpoint_id: UUID) -> None: """Evict cache entries related to an endpoint. Called on revision changes or probe failures to force re-resolution on the next request. """ self._cache.invalidate(f"endpoint:{endpoint_id}") # Also clear all binding caches since they may reference this endpoint # We clear all bindings because we can't efficiently know which bindings # use this endpoint without scanning self._cache.invalidate("binding:") logger.info("Invalidated cache for endpoint %s", endpoint_id) def invalidate_all(self) -> None: """Full cache clear.""" self._cache.clear() logger.info("Full registry cache invalidated")