File size: 2,136 Bytes
3e363e2 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | """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
|