Coverage for agentos/tools/metrics.py: 24%

1""" 

2Metrics Collector for AgentOS. 

3 

4Lightweight in-process metrics with Counter, Gauge, Histogram, and Timer. 

5Thread-safe, zero external dependencies, Prometheus-style text exposition. 

6""" 

7 

8import threading 

9import time 

10from collections.abc import Callable 

11from typing import Any 

12 

13# ============================================================================ 

14# Histogram buckets (pre-defined for common latency ranges) 

15# ============================================================================ 

16 

17DEFAULT_BUCKETS = ( 

18 0.005, 

19 0.01, 

20 0.025, 

21 0.05, 

22 0.1, 

23 0.25, 

24 0.5, 

25 1.0, 

26 2.5, 

27 5.0, 

28 10.0, 

29) 

30LARGE_BUCKETS = ( 

31 0.1, 

32 0.5, 

33 1.0, 

34 5.0, 

35 10.0, 

36 30.0, 

37 60.0, 

38 300.0, 

39 600.0, 

40 1800.0, 

41 3600.0, 

42) 

43 

44 

45# ============================================================================ 

46# Metric types 

47# ============================================================================ 

48 

49 

50class Counter: 

51 """Monotonically increasing counter. Thread-safe.""" 

52 

53 def __init__(self, name: str, help_text: str = "", labels: dict[str, str] | None = None): 

54 self.name = name 

55 self.help = help_text 

56 self.labels = labels or {} 

57 self._value: float = 0.0 

58 self._lock = threading.Lock() 

59 

60 def inc(self, delta: float = 1.0) -> None: 

61 with self._lock: 

62 self._value += delta 

63 

64 def get(self) -> float: 

65 with self._lock: 

66 return self._value 

67 

68 def _format_labels(self) -> str: 

69 if not self.labels: 

70 return "" 

71 return "{" + ",".join(f'{k}="{v}"' for k, v in self.labels.items()) + "}" 

72 

73 

74class Gauge: 

75 """Value that can go up and down. Thread-safe.""" 

76 

77 def __init__(self, name: str, help_text: str = "", labels: dict[str, str] | None = None): 

78 self.name = name 

79 self.help = help_text 

80 self.labels = labels or {} 

81 self._value: float = 0.0 

82 self._lock = threading.Lock() 

83 

84 def set(self, value: float) -> None: 

85 with self._lock: 

86 self._value = value 

87 

88 def inc(self, delta: float = 1.0) -> None: 

89 with self._lock: 

90 self._value += delta 

91 

92 def dec(self, delta: float = 1.0) -> None: 

93 with self._lock: 

94 self._value -= delta 

95 

96 def get(self) -> float: 

97 with self._lock: 

98 return self._value 

99 

100 def _format_labels(self) -> str: 

101 if not self.labels: 

102 return "" 

103 return "{" + ",".join(f'{k}="{v}"' for k, v in self.labels.items()) + "}" 

104 

105 

106class Histogram: 

107 """Bucketed histogram with sum and count. Thread-safe.""" 

108 

109 def __init__( 

110 self, 

111 name: str, 

112 help_text: str = "", 

113 labels: dict[str, str] | None = None, 

114 buckets: tuple = DEFAULT_BUCKETS, 

115 ): 

116 self.name = name 

117 self.help = help_text 

118 self.labels = labels or {} 

119 self.buckets = tuple(sorted(buckets)) 

120 self._lock = threading.Lock() 

121 self._bucket_counts: list[int] = [0] * (len(self.buckets) + 1) # +1 for +Inf 

122 self._sum: float = 0.0 

123 self._count: int = 0 

124 

125 def observe(self, value: float) -> None: 

126 with self._lock: 

127 self._sum += value 

128 self._count += 1 

129 for i, bound in enumerate(self.buckets): 

130 if value <= bound: 

131 self._bucket_counts[i] += 1 

132 return 

133 self._bucket_counts[-1] += 1 # +Inf bucket 

134 

135 def get(self) -> dict[str, Any]: 

136 with self._lock: 

137 return { 

138 "sum": self._sum, 

139 "count": self._count, 

140 "buckets": dict(zip(self.buckets + ("+Inf",), self._bucket_counts)), 

141 } 

142 

143 def p50(self) -> float: 

144 return self._percentile(0.50) 

145 

146 def p90(self) -> float: 

147 return self._percentile(0.90) 

148 

149 def p99(self) -> float: 

150 return self._percentile(0.99) 

151 

152 def _percentile(self, p: float) -> float: 

153 with self._lock: 

154 if self._count == 0: 

155 return 0.0 

156 target = int(self._count * p) 

157 accumulated = 0 

158 for i, bc in enumerate(self._bucket_counts): 

159 accumulated += bc 

160 if accumulated >= target: 

161 if i < len(self.buckets): 

162 return self.buckets[i] 

163 return self.buckets[-1] if self.buckets else 0.0 

164 return self.buckets[-1] if self.buckets else 0.0 

165 

166 def _format_labels(self) -> str: 

167 if not self.labels: 

168 return "" 

169 return "{" + ",".join(f'{k}="{v}"' for k, v in self.labels.items()) + "}" 

170 

171 

172class Timer: 

173 """Convenience wrapper: Histogram for timing. Also tracks rate via internal counter.""" 

174 

175 def __init__( 

176 self, 

177 name: str, 

178 help_text: str = "", 

179 labels: dict[str, str] | None = None, 

180 buckets: tuple = DEFAULT_BUCKETS, 

181 ): 

182 self.name = name 

183 self.help = help_text 

184 self.histogram = Histogram(name, help_text, labels, buckets) 

185 self._call_count = Counter(name + "_total", help_text, labels) 

186 

187 def time(self, fn: Callable[..., Any], *args, **kwargs) -> Any: 

188 start = time.perf_counter() 

189 try: 

190 return fn(*args, **kwargs) 

191 finally: 

192 elapsed = time.perf_counter() - start 

193 self.histogram.observe(elapsed) 

194 self._call_count.inc() 

195 

196 def __enter__(self): 

197 self._start = time.perf_counter() 

198 return self 

199 

200 def __exit__(self, *args): 

201 elapsed = time.perf_counter() - self._start 

202 self.histogram.observe(elapsed) 

203 self._call_count.inc() 

204 

205 def get(self) -> dict[str, Any]: 

206 return { 

207 "histogram": self.histogram.get(), 

208 "count": self._call_count.get(), 

209 } 

210 

211 

212# ============================================================================ 

213# MetricsCollector (Registry) 

214# ============================================================================ 

215 

216 

217class MetricsCollector: 

218 """Global registry for metrics. Provides Prometheus text format exposition.""" 

219 

220 def __init__(self, namespace: str = ""): 

221 self.namespace = namespace 

222 self._metrics: dict[str, Any] = {} 

223 self._lock = threading.Lock() 

224 

225 def _full_name(self, name: str) -> str: 

226 if self.namespace: 

227 return f"{self.namespace}_{name}" 

228 return name 

229 

230 def counter( 

231 self, name: str, help_text: str = "", labels: dict[str, str] | None = None 

232 ) -> Counter: 

233 full = self._full_name(name) 

234 with self._lock: 

235 if full not in self._metrics: 

236 self._metrics[full] = Counter(full, help_text, labels) 

237 return self._metrics[full] 

238 

239 def gauge(self, name: str, help_text: str = "", labels: dict[str, str] | None = None) -> Gauge: 

240 full = self._full_name(name) 

241 with self._lock: 

242 if full not in self._metrics: 

243 self._metrics[full] = Gauge(full, help_text, labels) 

244 return self._metrics[full] 

245 

246 def histogram( 

247 self, 

248 name: str, 

249 help_text: str = "", 

250 labels: dict[str, str] | None = None, 

251 buckets: tuple = DEFAULT_BUCKETS, 

252 ) -> Histogram: 

253 full = self._full_name(name) 

254 with self._lock: 

255 if full not in self._metrics: 

256 self._metrics[full] = Histogram(full, help_text, labels, buckets) 

257 return self._metrics[full] 

258 

259 def timer( 

260 self, 

261 name: str, 

262 help_text: str = "", 

263 labels: dict[str, str] | None = None, 

264 buckets: tuple = DEFAULT_BUCKETS, 

265 ) -> Timer: 

266 full = self._full_name(name) 

267 with self._lock: 

268 if full not in self._metrics: 

269 self._metrics[full] = Timer(full, help_text, labels, buckets) 

270 return self._metrics[full] 

271 

272 def get(self, name: str) -> Any | None: 

273 return self._metrics.get(self._full_name(name)) 

274 

275 def list_metrics(self) -> list[str]: 

276 with self._lock: 

277 return list(self._metrics.keys()) 

278 

279 def get_all(self) -> dict[str, Any]: 

280 """Return raw values for all metrics (for programmatic use).""" 

281 result = {} 

282 with self._lock: 

283 for name, metric in self._metrics.items(): 

284 if hasattr(metric, "get"): 

285 result[name] = metric.get() 

286 return result 

287 

288 def to_prometheus(self) -> str: 

289 """Export all metrics in Prometheus text format.""" 

290 lines = [] 

291 with self._lock: 

292 for name, metric in self._metrics.items(): 

293 if metric.help: 

294 lines.append(f"# HELP {name} {metric.help}") 

295 lines.append( 

296 f"# TYPE {name} histogram" 

297 if isinstance(metric, (Histogram, Timer)) 

298 else ( 

299 f"# TYPE {name} gauge" 

300 if isinstance(metric, Gauge) 

301 else f"# TYPE {name} counter" 

302 ) 

303 ) 

304 

305 if isinstance(metric, Counter): 

306 lbl = metric._format_labels() 

307 lines.append(f"{name}{lbl} {metric.get()}") 

308 elif isinstance(metric, Gauge): 

309 lbl = metric._format_labels() 

310 lines.append(f"{name}{lbl} {metric.get()}") 

311 elif isinstance(metric, Histogram): 

312 lbl = metric._format_labels() 

313 data = metric.get() 

314 lines.append(f"{name}_count{lbl} {data['count']}") 

315 lines.append(f"{name}_sum{lbl} {data['sum']}") 

316 for bucket_name in metric.buckets + ("+Inf",): 

317 bval = data["buckets"].get(bucket_name, 0) 

318 # Prometheus: bucket labels use 'le' key 

319 le_label = '{le="' + str(bucket_name) + '"}' 

320 lines.append(f"{name}_bucket{lbl}{le_label} {bval}") 

321 elif isinstance(metric, Timer): 

322 lbl = metric.histogram._format_labels() 

323 hdata = metric.histogram.get() 

324 lines.append(f"{name}_count{lbl} {hdata['count']}") 

325 lines.append(f"{name}_sum{lbl} {hdata['sum']}") 

326 for bucket_name in metric.histogram.buckets + ("+Inf",): 

327 bval = hdata["buckets"].get(bucket_name, 0) 

328 le_label = '{le="' + str(bucket_name) + '"}' 

329 lines.append(f"{name}_bucket{lbl}{le_label} {bval}") 

330 

331 return "\n".join(lines) + "\n" 

332 

333 

334# ============================================================================ 

335# Global singleton 

336# ============================================================================ 

337 

338_default_collector: MetricsCollector | None = None 

339_collector_lock = threading.Lock() 

340 

341 

342def get_metrics_collector(namespace: str = "") -> MetricsCollector: 

343 global _default_collector 

344 if _default_collector is None: 

345 with _collector_lock: 

346 if _default_collector is None: 

347 _default_collector = MetricsCollector(namespace=namespace or "agentos") 

348 return _default_collector