Spaces:
Running
Running
File size: 2,668 Bytes
d2258e5 | 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 67 68 69 70 71 72 73 | from __future__ import annotations
from dataclasses import dataclass
PREDICTOR_SCOPES = {"global_ema", "per_tool_ema"}
@dataclass
class _Estimate:
mean_s: float
count: int = 0
class OnlineToolGapPredictor:
"""Tiny online tool-duration predictor used by adaptive KV tiering.
The predictor deliberately stays transparent: it keeps an exponentially
weighted moving average (EWMA) of observed tool-call durations. In
``per_tool_ema`` mode each tool family gets its own estimate, backed by a
global EWMA until enough observations have arrived. No future tool duration
is visible at decision time.
"""
def __init__(
self,
*,
initial_mean_s: float,
alpha: float = 0.30,
min_observations: int = 2,
scope: str = "per_tool_ema",
) -> None:
if scope not in PREDICTOR_SCOPES:
raise ValueError(f"Unsupported predictor scope: {scope}")
self.scope = scope
self.alpha = min(max(float(alpha), 0.01), 1.0)
self.min_observations = max(int(min_observations), 1)
self.global_estimate = _Estimate(max(float(initial_mean_s), 0.0), 0)
self.per_tool: dict[str, _Estimate] = {}
@staticmethod
def _update(estimate: _Estimate, value_s: float, alpha: float) -> None:
value_s = max(float(value_s), 0.0)
if estimate.count == 0:
estimate.mean_s = value_s
else:
estimate.mean_s = alpha * value_s + (1.0 - alpha) * estimate.mean_s
estimate.count += 1
def predict(self, tool_kind: str) -> tuple[float, str, int]:
if self.scope == "per_tool_ema":
estimate = self.per_tool.get(tool_kind)
if estimate is not None and estimate.count >= self.min_observations:
return estimate.mean_s, "tool", estimate.count
return self.global_estimate.mean_s, "global", self.global_estimate.count
def observe(self, tool_kind: str, duration_s: float) -> None:
self._update(self.global_estimate, duration_s, self.alpha)
estimate = self.per_tool.setdefault(tool_kind, _Estimate(self.global_estimate.mean_s, 0))
self._update(estimate, duration_s, self.alpha)
def snapshot(self) -> dict:
return {
"scope": self.scope,
"alpha": self.alpha,
"min_observations": self.min_observations,
"global": {"mean_s": self.global_estimate.mean_s, "count": self.global_estimate.count},
"per_tool": {
key: {"mean_s": value.mean_s, "count": value.count}
for key, value in sorted(self.per_tool.items())
},
}
|