Coverage for agentos/core/observability.py: 0%

1"""AgentOS Observability — metrics, structured logging, and tracing.

3Production-grade observability module with:

4- Counters, Gauges, Histograms, Summaries (Prometheus-compatible naming)

5- Structured JSON logging with correlation IDs

6- Span-based tracing for request lifecycle

7- Export registry for Prometheus scraping

8- Zero external deps beyond stdlib

10Design: ~350 lines. Async-safe, thread-safe.

11"""

13from __future__ import annotations

15import json

16import logging

17import threading

18import time

19from collections import OrderedDict

20from contextvars import ContextVar

21from dataclasses import dataclass, field

22from enum import StrEnum

23from typing import Any

24from uuid import uuid4

26# ============================================================================

27# Correlation & context

28# ============================================================================

30correlation_id: ContextVar[str] = ContextVar("correlation_id", default="")

31span_id: ContextVar[str] = ContextVar("span_id", default="")

34def new_correlation_id() -> str:

35 cid = str(uuid4())[:12]

36 correlation_id.set(cid)

37 return cid

40def get_correlation_id() -> str:

41 return correlation_id.get()

44# ============================================================================

45# Metrics Core

46# ============================================================================

49class MetricType(StrEnum):

50 COUNTER = "counter"

51 GAUGE = "gauge"

52 HISTOGRAM = "histogram"

53 SUMMARY = "summary"

56@dataclass

57class MetricLabel:

58 name: str

59 value: str

62class MetricValue:

63 """Thread-safe atomic metric value for counters and gauges."""

65 def __init__(self, initial: float = 0.0):

66 self._value = initial

67 self._lock = threading.Lock()

69 def get(self) -> float:

70 with self._lock:

71 return self._value

73 def set(self, v: float) -> None:

74 with self._lock:

75 self._value = v

77 def add(self, delta: float) -> None:

78 with self._lock:

79 self._value += delta

81 def inc(self) -> None:

82 self.add(1.0)

85class _HistogramBucket:

86 """Thread-safe histogram bucket."""

88 def __init__(self):

89 self._values: list[float] = []

90 self._lock = threading.Lock()

91 self._sum = 0.0

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

94 with self._lock:

95 self._values.append(value)

96 self._sum += value

98 def snapshot(self) -> dict[str, Any]:

99 with self._lock:

100 vals = sorted(self._values) if self._values else []

101 count = len(vals)

102 return {

103 "count": count,

104 "sum": round(self._sum, 6),

105 "min": vals[0] if count else 0,

106 "max": vals[-1] if count else 0,

107 "avg": round(self._sum / count, 6) if count else 0,

108 "p50": _percentile(vals, 50),

109 "p90": _percentile(vals, 90),

110 "p95": _percentile(vals, 95),

111 "p99": _percentile(vals, 99),

112 }

113

114

115def _percentile(sorted_vals: list[float], p: int) -> float:

116 if not sorted_vals:

117 return 0.0

118 k = (len(sorted_vals) - 1) * p / 100.0

119 f = int(k)

120 c = k - f

121 if f + 1 < len(sorted_vals):

122 return sorted_vals[f] + c * (sorted_vals[f + 1] - sorted_vals[f])

123 return sorted_vals[f]

124

125

126# ============================================================================

127# Metrics

128# ============================================================================

129

130

131class Counter:

132 """Monotonic counter (only increases). Prometheus-compatible."""

133

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

135 self.name = name

136 self.help = help_text

137 self.labels = labels or {}

138 self._value = MetricValue(0.0)

139

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

141 self._value.add(delta)

142

143 def value(self) -> float:

144 return self._value.get()

145

146

147class Gauge:

148 """Gauge that can go up and down."""

149

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

151 self.name = name

152 self.help = help_text

153 self.labels = labels or {}

154 self._value = MetricValue(0.0)

155

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

157 self._value.set(value)

158

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

160 self._value.add(delta)

161

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

163 self._value.add(-delta)

164

165 def value(self) -> float:

166 return self._value.get()

167

168

169class Histogram:

170 """Histogram for latency/distribution metrics."""

171

172 def __init__(

173 self,

174 name: str,

175 help_text: str = "",

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

177 buckets: list[float] | None = None,

178 ):

179 self.name = name

180 self.help = help_text

181 self.labels = labels or {}

182 self.buckets = buckets or [0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0]

183 self._bucket = _HistogramBucket()

184

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

186 self._bucket.observe(value)

187

188 def snapshot(self) -> dict[str, Any]:

189 return self._bucket.snapshot()

190

191

192class Summary(Histogram):

193 """Summary metric (alias for Histogram with quantiles)."""

194

195

196

197# ============================================================================

198# Metrics Registry

199# ============================================================================

200

201

202class MetricsRegistry:

203 """Collect, query, and export all registered metrics."""

204

205 def __init__(self):

206 self._metrics: dict[str, Any] = OrderedDict()

207 self._lock = threading.Lock()

208

209 def counter(

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

211 ) -> Counter:

212 with self._lock:

213 if name not in self._metrics:

214 self._metrics[name] = Counter(name=name, help_text=help_text, labels=labels)

215 return self._metrics[name]

216

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

218 with self._lock:

219 if name not in self._metrics:

220 self._metrics[name] = Gauge(name=name, help_text=help_text, labels=labels)

221 return self._metrics[name]

222

223 def histogram(

224 self,

225 name: str,

226 help_text: str = "",

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

228 buckets: list[float] | None = None,

229 ) -> Histogram:

230 with self._lock:

231 if name not in self._metrics:

232 self._metrics[name] = Histogram(

233 name=name, help_text=help_text, labels=labels, buckets=buckets

234 )

235 return self._metrics[name]

236

237 def get(self, name: str):

238 return self._metrics.get(name)

239

240 def export_prometheus(self) -> str:

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

242 lines: list[str] = []

243 with self._lock:

244 for metric in self._metrics.values():

245 if metric.help:

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

247 lines.append(f"# TYPE {metric.name} {self._prometheus_type(metric)}")

248

249 label_str = self._format_labels(metric.labels)

250 if isinstance(metric, (Counter, Gauge)):

251 lines.append(f"{metric.name}{label_str} {metric.value()}")

252 elif isinstance(metric, Histogram):

253 snap = metric.snapshot()

254 lines.append(f"{metric.name}_count{label_str} {snap['count']}")

255 lines.append(f"{metric.name}_sum{label_str} {snap['sum']}")

256 # Bucket labels

257 for b in metric.buckets:

258 lines.append(f"{metric.name}_bucket{{le=\"{b}\"}} {snap['count']}")

259 lines.append(f"{metric.name}_bucket{{le=\"+Inf\"}} {snap['count']}")

260 lines.append("")

261 return "\n".join(lines)

262

263 def export_dict(self) -> dict[str, Any]:

264 """Export all metrics as a Python dict (for JSON APIs)."""

265 result: dict[str, Any] = {}

266 with self._lock:

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

268 if isinstance(metric, (Counter, Gauge)):

269 result[name] = metric.value()

270 elif isinstance(metric, Histogram):

271 result[name] = metric.snapshot()

272 return result

273

274 def _prometheus_type(self, metric) -> str:

275 if isinstance(metric, Counter):

276 return "counter"

277 if isinstance(metric, Gauge):

278 return "gauge"

279 if isinstance(metric, Histogram):

280 return "histogram"

281 return "untyped"

282

283 def _format_labels(self, labels: dict[str, str]) -> str:

284 if not labels:

285 return ""

286 parts = [f'{k}="{v}"' for k, v in labels.items()]

287 return "{" + ",".join(parts) + "}"

288

289

290# ============================================================================

291# Tracing (minimal span-based)

292# ============================================================================

293

294

295@dataclass

296class SpanContext:

297 trace_id: str = field(default_factory=lambda: str(uuid4())[:16])

298 span_id: str = field(default_factory=lambda: str(uuid4())[:12])

299 parent_span_id: str | None = None

300

301

302class Span:

303 """Minimal span for request tracing."""

304

305 def __init__(

306 self,

307 name: str,

308 parent: Span | None = None,

309 trace_id: str | None = None,

310 ):

311 self.name = name

312 self.context = SpanContext(

313 trace_id=trace_id or (parent.context.trace_id if parent else str(uuid4())[:16]),

314 span_id=str(uuid4())[:12],

315 parent_span_id=parent.context.span_id if parent else None,

316 )

317 self.start_time = time.monotonic()

318 self.end_time: float | None = None

319 self._tags: dict[str, str] = {}

320 self._events: list[dict[str, Any]] = []

321

322 def set_tag(self, key: str, value: str) -> None:

323 self._tags[key] = value

324

325 def add_event(self, name: str, **attributes) -> None:

326 self._events.append(

327 {

328 "name": name,

329 "timestamp": time.monotonic(),

330 "attributes": attributes,

331 }

332 )

333

334 def finish(self) -> None:

335 self.end_time = time.monotonic()

336

337 @property

338 def duration_ms(self) -> float:

339 end = self.end_time or time.monotonic()

340 return (end - self.start_time) * 1000

341

342 def to_dict(self) -> dict[str, Any]:

343 return {

344 "name": self.name,

345 "trace_id": self.context.trace_id,

346 "span_id": self.context.span_id,

347 "parent_span_id": self.context.parent_span_id,

348 "duration_ms": round(self.duration_ms, 3),

349 "tags": self._tags,

350 "events": self._events,

351 }

352

353

354class Tracer:

355 """Creates spans and manages active trace context."""

356

357 _active_span: ContextVar[Span | None] = ContextVar("active_span", default=None)

358

359 @classmethod

360 def noop(cls) -> Tracer:

361 """Return a no-op tracer that creates no spans."""

362 return cls()

363

364 def start_span(self, name: str) -> Span:

365 parent = self._active_span.get()

366 span = Span(name=name, parent=parent)

367 self._active_span.set(span)

368 return span

369

370 def end_span(self, span: Span) -> None:

371 span.finish()

372 # Restore parent span if exists

373 if span.context.parent_span_id is not None:

374 # Simple case: restore parent

375 pass

376 self._active_span.set(None)

377

378 @property

379 def active_span(self) -> Span | None:

380 return self._active_span.get()

381

382 def span_context(self) -> SpanContext:

383 span = self._active_span.get()

384 if span:

385 return span.context

386 return SpanContext()

387

388

389# ============================================================================

390# Structured Logging

391# ============================================================================

392

393

394class JsonFormatter(logging.Formatter):

395 """JSON structured log formatter with correlation ID injection."""

396

397 def format(self, record: logging.LogRecord) -> str:

398 log_entry: dict[str, Any] = OrderedDict()

399 log_entry["timestamp"] = self.formatTime(record, self.datefmt)

400 log_entry["level"] = record.levelname

401 log_entry["logger"] = record.name

402 log_entry["message"] = record.getMessage()

403

404 cid = correlation_id.get()

405 if cid:

406 log_entry["correlation_id"] = cid

407

408 sid = span_id.get()

409 if sid:

410 log_entry["span_id"] = sid

411

412 if record.exc_info and record.exc_info[0]:

413 log_entry["exception"] = self.formatException(record.exc_info)

414

415 if hasattr(record, "extra_fields"):

416 log_entry.update(getattr(record, "extra_fields", {}))

417

418 return json.dumps(log_entry, default=str, ensure_ascii=False)

419

420

421def setup_structured_logging(level: int = logging.INFO):

422 """Configure root logger with JSON structured output."""

423 handler = logging.StreamHandler()

424 handler.setFormatter(JsonFormatter())

425 root = logging.getLogger()

426 root.handlers.clear()

427 root.addHandler(handler)

428 root.setLevel(level)

429

430

431# ============================================================================

432# Pre-built metric sets

433# ============================================================================

434

435

436@dataclass

437class StandardMetrics:

438 """Common agent metrics used across AgentOS."""

439

440 def __init__(self, registry: MetricsRegistry, prefix: str = "agentos"):

441 self.registry = registry

442 self.p = prefix

443

444 # Request metrics

445 self.request_count = registry.counter(

446 f"{prefix}_requests_total", "Total requests processed"

447 )

448 self.request_duration = registry.histogram(

449 f"{prefix}_request_duration_seconds", "Request duration in seconds"

450 )

451 self.request_errors = registry.counter(

452 f"{prefix}_request_errors_total", "Total request errors"

453 )

454

455 # Agent metrics

456 self.active_agents = registry.gauge(f"{prefix}_active_agents", "Currently running agents")

457 self.agent_invocations = registry.counter(

458 f"{prefix}_agent_invocations_total", "Total agent invocations"

459 )

460

461 # Resource metrics

462 self.memory_bytes = registry.gauge(

463 f"{prefix}_memory_bytes", "Current memory usage in bytes"

464 )

465

466

467# ============================================================================

468# Module-level instances

469# ============================================================================

470

471default_registry = MetricsRegistry()

472default_tracer = Tracer()

473default_metrics = StandardMetrics(default_registry)