Multi-stage evidence-grounded inference architecture replacing the monolithic 9B model extraction pipeline. CPU-first specialist services handle routine extraction while the 9B vLLM model is preserved for semantic adjudication of ambiguous cases. Key components: - Capability-aware inference gateway (OpenAI-compatible + Ollama) - Endpoint registry with DB migrations and REST API - Sentence-aware document segmenter (property tests) - Deterministic financial parsing with offset integrity - Symbol resolution with ambiguity detection - Specialist service (GLiNER2, dynamic batching, K8s deployment) - Company-specific sentiment (FinBERT, calibration) - Retrieval-based novelty and duplicate detection - Confidence calibration pipeline - Deterministic routing engine (property tests) - 9B adjudication layer with VRAM gating - Stock-specific impact model (features, labels, baseline, trained) - Pipeline orchestrator (state machine, queues, leases, feature flags) - Bounded parallelism (async workers, semaphore, load shedding) - Observability (tracing, metrics, alerts) - Compatibility adapter (v3→v2 golden mapping tests) - Shadow/canary promotion framework - Active learning and fine-tuning pipeline Test results: 1,161 tests pass, ruff lint clean. All 282 spec tasks completed.
320 lines
11 KiB
Python
320 lines
11 KiB
Python
"""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")
|