| """Production runtime wrapper with graph/model/version-aware caching.""" |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| from collections.abc import Mapping |
| from dataclasses import dataclass, replace |
|
|
| import torch |
|
|
| from graph.graph import GraphStore |
| from graph.model import GraphModel |
| from graph.training.calibration import TemperatureCalibrator |
| from graph.types import GraphQuery, ReasoningResult |
|
|
|
|
| @dataclass(frozen=True) |
| class CacheKey: |
| tenant_id: str |
| graph_version: str |
| model_version: str |
| query_hash: str |
|
|
|
|
| class ReasoningRuntime: |
| def __init__( |
| self, |
| model: GraphModel, |
| graphs: Mapping[str, GraphStore] | None = None, |
| calibrator: TemperatureCalibrator | None = None, |
| ) -> None: |
| self.model = model |
| self.graphs = dict(graphs or {}) |
| self.calibrator = calibrator |
| self.cache: dict[CacheKey, ReasoningResult] = {} |
|
|
| def register_graph(self, tenant_id: str, graph: GraphStore) -> None: |
| self.graphs[tenant_id] = graph |
| for key in tuple(self.cache): |
| if key.tenant_id == tenant_id: |
| del self.cache[key] |
|
|
| def reason( |
| self, tenant_id: str, query: GraphQuery, *, use_cache: bool = True |
| ) -> ReasoningResult: |
| graph = self.graphs.get(tenant_id) |
| if graph is None: |
| raise KeyError(f"unknown tenant graph {tenant_id!r}") |
| key = CacheKey( |
| tenant_id, |
| graph.graph_version, |
| self.model.config.model_version, |
| hashlib.sha256(query.canonical_text().encode()).hexdigest(), |
| ) |
| if use_cache and key in self.cache: |
| return self.cache[key] |
| result = self.model.reason(graph, query) |
| if self.calibrator is not None: |
| probability = torch.tensor(result.confidence).clamp(1e-6, 1 - 1e-6) |
| logit = torch.log(probability / (1 - probability)) |
| calibrated = float(self.calibrator.transform(logit).item()) |
| result = replace(result, confidence=calibrated, calibrated=True) |
| if use_cache: |
| self.cache[key] = result |
| return result |
|
|
