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

197 statements  

« prev     ^ index     » next       coverage.py v7.14.3, created at 2026-07-06 10:59 +0800

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 typing import Any, Callable, Dict, List, Optional 

11 

12 

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

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

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

16 

17DEFAULT_BUCKETS = ( 

18 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 

19) 

20LARGE_BUCKETS = ( 

21 0.1, 0.5, 1.0, 5.0, 10.0, 30.0, 60.0, 300.0, 600.0, 1800.0, 3600.0, 

22) 

23 

24 

25# ============================================================================ 

26# Metric types 

27# ============================================================================ 

28 

29class Counter: 

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

31 

32 def __init__(self, name: str, help_text: str = "", labels: Optional[Dict[str, str]] = None): 

33 self.name = name 

34 self.help = help_text 

35 self.labels = labels or {} 

36 self._value: float = 0.0 

37 self._lock = threading.Lock() 

38 

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

40 with self._lock: 

41 self._value += delta 

42 

43 def get(self) -> float: 

44 with self._lock: 

45 return self._value 

46 

47 def _format_labels(self) -> str: 

48 if not self.labels: 

49 return "" 

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

51 

52 

53class Gauge: 

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

55 

56 def __init__(self, name: str, help_text: str = "", labels: Optional[Dict[str, str]] = None): 

57 self.name = name 

58 self.help = help_text 

59 self.labels = labels or {} 

60 self._value: float = 0.0 

61 self._lock = threading.Lock() 

62 

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

64 with self._lock: 

65 self._value = value 

66 

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

68 with self._lock: 

69 self._value += delta 

70 

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

72 with self._lock: 

73 self._value -= delta 

74 

75 def get(self) -> float: 

76 with self._lock: 

77 return self._value 

78 

79 def _format_labels(self) -> str: 

80 if not self.labels: 

81 return "" 

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

83 

84 

85class Histogram: 

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

87 

88 def __init__( 

89 self, 

90 name: str, 

91 help_text: str = "", 

92 labels: Optional[Dict[str, str]] = None, 

93 buckets: tuple = DEFAULT_BUCKETS, 

94 ): 

95 self.name = name 

96 self.help = help_text 

97 self.labels = labels or {} 

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

99 self._lock = threading.Lock() 

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

101 self._sum: float = 0.0 

102 self._count: int = 0 

103 

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

105 with self._lock: 

106 self._sum += value 

107 self._count += 1 

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

109 if value <= bound: 

110 self._bucket_counts[i] += 1 

111 return 

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

113 

114 def get(self) -> Dict[str, Any]: 

115 with self._lock: 

116 return { 

117 "sum": self._sum, 

118 "count": self._count, 

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

120 } 

121 

122 def p50(self) -> float: 

123 return self._percentile(0.50) 

124 

125 def p90(self) -> float: 

126 return self._percentile(0.90) 

127 

128 def p99(self) -> float: 

129 return self._percentile(0.99) 

130 

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

132 with self._lock: 

133 if self._count == 0: 

134 return 0.0 

135 target = int(self._count * p) 

136 accumulated = 0 

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

138 accumulated += bc 

139 if accumulated >= target: 

140 if i < len(self.buckets): 

141 return self.buckets[i] 

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

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

144 

145 def _format_labels(self) -> str: 

146 if not self.labels: 

147 return "" 

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

149 

150 

151class Timer: 

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

153 

154 def __init__( 

155 self, 

156 name: str, 

157 help_text: str = "", 

158 labels: Optional[Dict[str, str]] = None, 

159 buckets: tuple = DEFAULT_BUCKETS, 

160 ): 

161 self.name = name 

162 self.help = help_text 

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

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

165 

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

167 start = time.perf_counter() 

168 try: 

169 return fn(*args, **kwargs) 

170 finally: 

171 elapsed = time.perf_counter() - start 

172 self.histogram.observe(elapsed) 

173 self._call_count.inc() 

174 

175 def __enter__(self): 

176 self._start = time.perf_counter() 

177 return self 

178 

179 def __exit__(self, *args): 

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

181 self.histogram.observe(elapsed) 

182 self._call_count.inc() 

183 

184 def get(self) -> Dict[str, Any]: 

185 return { 

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

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

188 } 

189 

190 

191# ============================================================================ 

192# MetricsCollector (Registry) 

193# ============================================================================ 

194 

195class MetricsCollector: 

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

197 

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

199 self.namespace = namespace 

200 self._metrics: Dict[str, Any] = {} 

201 self._lock = threading.Lock() 

202 

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

204 if self.namespace: 

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

206 return name 

207 

208 def counter(self, name: str, help_text: str = "", labels: Optional[Dict[str, str]] = None) -> Counter: 

209 full = self._full_name(name) 

210 with self._lock: 

211 if full not in self._metrics: 

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

213 return self._metrics[full] 

214 

215 def gauge(self, name: str, help_text: str = "", labels: Optional[Dict[str, str]] = None) -> Gauge: 

216 full = self._full_name(name) 

217 with self._lock: 

218 if full not in self._metrics: 

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

220 return self._metrics[full] 

221 

222 def histogram( 

223 self, 

224 name: str, 

225 help_text: str = "", 

226 labels: Optional[Dict[str, str]] = None, 

227 buckets: tuple = DEFAULT_BUCKETS, 

228 ) -> Histogram: 

229 full = self._full_name(name) 

230 with self._lock: 

231 if full not in self._metrics: 

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

233 return self._metrics[full] 

234 

235 def timer( 

236 self, 

237 name: str, 

238 help_text: str = "", 

239 labels: Optional[Dict[str, str]] = None, 

240 buckets: tuple = DEFAULT_BUCKETS, 

241 ) -> Timer: 

242 full = self._full_name(name) 

243 with self._lock: 

244 if full not in self._metrics: 

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

246 return self._metrics[full] 

247 

248 def get(self, name: str) -> Optional[Any]: 

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

250 

251 def list_metrics(self) -> List[str]: 

252 with self._lock: 

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

254 

255 def get_all(self) -> Dict[str, Any]: 

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

257 result = {} 

258 with self._lock: 

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

260 if hasattr(metric, "get"): 

261 result[name] = metric.get() 

262 return result 

263 

264 def to_prometheus(self) -> str: 

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

266 lines = [] 

267 with self._lock: 

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

269 if metric.help: 

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

271 lines.append(f"# TYPE {name} histogram" if isinstance(metric, (Histogram, Timer)) else 

272 f"# TYPE {name} gauge" if isinstance(metric, Gauge) else 

273 f"# TYPE {name} counter") 

274 

275 if isinstance(metric, Counter): 

276 lbl = metric._format_labels() 

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

278 elif isinstance(metric, Gauge): 

279 lbl = metric._format_labels() 

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

281 elif isinstance(metric, Histogram): 

282 lbl = metric._format_labels() 

283 data = metric.get() 

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

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

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

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

288 # Prometheus: bucket labels use 'le' key 

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

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

291 elif isinstance(metric, Timer): 

292 lbl = metric.histogram._format_labels() 

293 hdata = metric.histogram.get() 

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

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

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

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

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

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

300 

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

302 

303 

304# ============================================================================ 

305# Global singleton 

306# ============================================================================ 

307 

308_default_collector: Optional[MetricsCollector] = None 

309_collector_lock = threading.Lock() 

310 

311 

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

313 global _default_collector 

314 if _default_collector is None: 

315 with _collector_lock: 

316 if _default_collector is None: 

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

318 return _default_collector