Coverage for agentos/observability/__init_

1"""

2AgentOS v1.14.3 — Observability Platform (OpenTelemetry Integration).

4Production-grade observability for AgentOS agent pipelines:

5- Distributed tracing (OpenTelemetry spans)

6- Metrics collection (Prometheus-compatible counters, histograms, gauges)

7- Structured logging with correlation IDs

8- Agent lifecycle events instrumentation

9- Cost tracking dashboard

10- Latency breakdown by pipeline stage

12Architecture:

13 AgentOS Pipeline

14 ├── OTel Tracer (auto-instrumented)

15 │ ├── Agent invocation span

16 │ │ ├── LLM call span

17 │ │ ├── Tool call span

18 │ │ └── Memory retrieval span

19 ├── MetricsExporter (Prometheus)

20 └── StructuredLogger (JSON)

22Inspired by: LangSmith, Weave, Arize Phoenix

23"""

25from __future__ import annotations

27import time

28import uuid

29from collections import defaultdict

30from collections.abc import Callable, Iterator

31from contextlib import contextmanager

32from dataclasses import dataclass, field

33from enum import StrEnum

34from typing import (

35 Any,

36)

38# ── Span & Trace Types ──────────────────────

41class SpanKind(StrEnum):

42 AGENT = "agent"

43 LLM = "llm"

44 TOOL = "tool"

45 MEMORY = "memory"

46 RETRIEVAL = "retrieval"

47 CHAIN = "chain"

48 EMBEDDING = "embedding"

51class SpanStatus(StrEnum):

52 OK = "ok"

53 ERROR = "error"

54 TIMEOUT = "timeout"

57@dataclass

58class SpanEvent:

59 """Span 中的事件。"""

61 name: str

62 timestamp: float = field(default_factory=time.time)

63 attributes: dict[str, Any] = field(default_factory=dict)

66@dataclass

67class Span:

68 """OpenTelemetry 风格的 Span。"""

70 span_id: str = field(default_factory=lambda: f"span-{uuid.uuid4().hex[:12]}")

71 parent_id: str | None = None

72 trace_id: str = field(default_factory=lambda: f"trace-{uuid.uuid4().hex[:16]}")

73 name: str = ""

74 kind: SpanKind = SpanKind.AGENT

76 start_time: float = field(default_factory=time.time)

77 end_time: float | None = None

78 status: SpanStatus = SpanStatus.OK

80 attributes: dict[str, Any] = field(default_factory=dict)

81 events: list[SpanEvent] = field(default_factory=list)

82 children: list[Span] = field(default_factory=list)

84 # Agent-specific

85 model_name: str = ""

86 input_tokens: int = 0

87 output_tokens: int = 0

88 cost_usd: float = 0.0

90 @property

91 def duration_ms(self) -> float:

92 end = self.end_time or time.time()

93 return (end - self.start_time) * 1000

95 def add_event(self, name: str, attributes: dict[str, Any] | None = None) -> SpanEvent:

96 event = SpanEvent(name=name, attributes=attributes or {})

97 self.events.append(event)

98 return event

100 def set_error(self, error: str) -> None:

101 self.status = SpanStatus.ERROR

102 self.attributes["error"] = error

103

104 def finish(self) -> None:

105 self.end_time = time.time()

106

107 def to_dict(self) -> dict:

108 return {

109 "span_id": self.span_id,

110 "parent_id": self.parent_id,

111 "trace_id": self.trace_id,

112 "name": self.name,

113 "kind": self.kind.value,

114 "duration_ms": self.duration_ms,

115 "status": self.status.value,

116 "attributes": self.attributes,

117 "model_name": self.model_name,

118 "input_tokens": self.input_tokens,

119 "output_tokens": self.output_tokens,

120 "cost_usd": self.cost_usd,

121 }

122

123

124@dataclass

125class Trace:

126 """一次完整的 Agent 调用链路。"""

127

128 trace_id: str = field(default_factory=lambda: f"trace-{uuid.uuid4().hex[:16]}")

129 root_span: Span | None = None

130 spans: list[Span] = field(default_factory=list)

131

132 @property

133 def total_duration_ms(self) -> float:

134 if not self.root_span:

135 return 0.0

136 return self.root_span.duration_ms

137

138 @property

139 def total_cost_usd(self) -> float:

140 return sum(s.cost_usd for s in self.spans)

141

142 @property

143 def total_tokens(self) -> int:

144 return sum(s.input_tokens + s.output_tokens for s in self.spans)

145

146 def to_dict(self) -> dict:

147 return {

148 "trace_id": self.trace_id,

149 "total_duration_ms": self.total_duration_ms,

150 "total_cost_usd": self.total_cost_usd,

151 "total_tokens": self.total_tokens,

152 "span_count": len(self.spans),

153 "spans": [s.to_dict() for s in self.spans],

154 }

155

156

157# ── Tracer ──────────────────────────────────

158

159

160class Tracer:

161 """AgentOS 追踪器。

162

163 自动检测 Agent/LMM/Tool/Memory 调用并创建 Span。

164

165 Usage:

166 tracer = Tracer()

167 with tracer.start_span("my_agent", SpanKind.AGENT) as span:

168 span.set_attribute("user_id", "123")

169 # ... agent logic ...

170 """

171

172 def __init__(self, service_name: str = "agentos"):

173 self._service_name = service_name

174 self._active_trace: Trace | None = None

175 self._span_stack: list[Span] = []

176 self._exporters: list[Callable] = []

177 self._trace_count: int = 0

178

179 def add_exporter(self, exporter: Callable[[Trace], None]) -> None:

180 """添加导出器。"""

181 self._exporters.append(exporter)

182

183 @contextmanager

184 def start_trace(self, name: str = "") -> Iterator[Trace]:

185 """创建新 Trace。"""

186 trace = Trace()

187 trace.trace_id = f"trace-{uuid.uuid4().hex[:16]}"

188

189 old_trace = self._active_trace

190 self._active_trace = trace

191

192 try:

193 yield trace

194 finally:

195 trace.root_span = trace.spans[0] if trace.spans else None

196 self._active_trace = old_trace

197 self._trace_count += 1

198

199 # Export

200 for exporter in self._exporters:

201 try:

202 exporter(trace)

203 except Exception:

204 pass

205

206 @contextmanager

207 def start_span(

208 self,

209 name: str,

210 kind: SpanKind = SpanKind.AGENT,

211 attributes: dict[str, Any] | None = None,

212 ) -> Iterator[Span]:

213 """在当前 Trace 中创建 Span。"""

214 span = Span(

215 name=name,

216 kind=kind,

217 trace_id=self._active_trace.trace_id if self._active_trace else "",

218 )

219

220 if self._span_stack:

221 span.parent_id = self._span_stack[-1].span_id

222 self._span_stack[-1].children.append(span)

223

224 if attributes:

225 span.attributes.update(attributes)

226

227 self._span_stack.append(span)

228

229 if self._active_trace:

230 self._active_trace.spans.append(span)

231

232 try:

233 yield span

234 except Exception as e:

235 span.set_error(str(e))

236 raise

237 finally:

238 span.finish()

239 self._span_stack.pop()

240

241 def record_llm_call(

242 self,

243 model: str,

244 input_tokens: int,

245 output_tokens: int,

246 cost_usd: float,

247 duration_ms: float,

248 ) -> None:

249 """记录 LLM 调用指标。"""

250 if self._span_stack:

251 span = self._span_stack[-1]

252 span.model_name = model

253 span.input_tokens = input_tokens

254 span.output_tokens = output_tokens

255 span.cost_usd = cost_usd

256

257 def record_tool_call(self, tool_name: str, success: bool, duration_ms: float) -> None:

258 """记录工具调用。"""

259 if self._span_stack:

260 span = self._span_stack[-1]

261 span.add_event(

262 "tool_call",

263 {

264 "tool_name": tool_name,

265 "success": success,

266 "duration_ms": duration_ms,

267 },

268 )

269

270

271# ── Metrics ─────────────────────────────────

272

273

274class MetricType(StrEnum):

275 COUNTER = "counter"

276 GAUGE = "gauge"

277 HISTOGRAM = "histogram"

278

279

280@dataclass

281class Metric:

282 """单个指标。"""

283

284 name: str

285 type: MetricType

286 description: str = ""

287 labels: dict[str, str] = field(default_factory=dict)

288 value: float = 0.0

289 timestamp: float = field(default_factory=time.time)

290

291

292class MetricsRegistry:

293 """Prometheus 风格的指标注册表。

294

295 Usage:

296 registry = MetricsRegistry()

297 counter = registry.counter("agent_invocations_total", "Total agent calls")

298 counter.inc()

299

300 histogram = registry.histogram("llm_latency_ms", "LLM call latency")

301 histogram.observe(1234.5)

302 """

303

304 def __init__(self):

305 self._metrics: dict[str, Metric] = {}

306 self._counters: dict[str, Counter] = {}

307 self._histograms: dict[str, Histogram] = {}

308 self._gauges: dict[str, Gauge] = {}

309

310 def counter(self, name: str, description: str = "") -> Counter:

311 if name not in self._counters:

312 c = Counter(name, description)

313 self._counters[name] = c

314 self._metrics[name] = Metric(

315 name=name, type=MetricType.COUNTER, description=description

316 )

317 return self._counters[name]

318

319 def histogram(

320 self, name: str, description: str = "", buckets: list[float] | None = None

321 ) -> Histogram:

322 if name not in self._histograms:

323 h = Histogram(name, description, buckets)

324 self._histograms[name] = h

325 self._metrics[name] = Metric(

326 name=name, type=MetricType.HISTOGRAM, description=description

327 )

328 return self._histograms[name]

329

330 def gauge(self, name: str, description: str = "") -> Gauge:

331 if name not in self._gauges:

332 g = Gauge(name, description)

333 self._gauges[name] = g

334 self._metrics[name] = Metric(name=name, type=MetricType.GAUGE, description=description)

335 return self._gauges[name]

336

337 def collect(self) -> list[dict]:

338 """收集所有指标的当前值（Prometheus scrape 格式）。"""

339 results = []

340 now = time.time()

341

342 for name, counter in self._counters.items():

343 results.append(

344 {

345 "name": name,

346 "type": "counter",

347 "value": counter.value,

348 "timestamp": now,

349 }

350 )

351

352 for name, histogram in self._histograms.items():

353 results.append(

354 {

355 "name": name,

356 "type": "histogram",

357 "count": histogram.count,

358 "sum": histogram.sum,

359 "buckets": dict(histogram.buckets),

360 "timestamp": now,

361 }

362 )

363

364 for name, gauge in self._gauges.items():

365 results.append(

366 {

367 "name": name,

368 "type": "gauge",

369 "value": gauge.value,

370 "timestamp": now,

371 }

372 )

373

374 return results

375

376 def to_prometheus_text(self) -> str:

377 """导出为 Prometheus 文本格式。"""

378 lines = []

379 for metric in self.collect():

380 lines.append(f"# HELP {metric['name']} AgentOS metric")

381 lines.append(f"# TYPE {metric['name']} {metric['type']}")

382

383 if metric["type"] == "histogram":

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

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

386 for bucket, count in metric.get("buckets", {}).items():

387 lines.append(f"{metric['name']}_bucket{{le=\"{bucket}\"}} {count}")

388 else:

389 lines.append(f"{metric['name']} {metric['value']}")

390

391 return "\n".join(lines)

392

393

394class Counter:

395 def __init__(self, name: str, description: str = ""):

396 self.name = name

397 self.description = description

398 self._value: float = 0.0

399

400 def inc(self, amount: float = 1.0) -> None:

401 self._value += amount

402

403 @property

404 def value(self) -> float:

405 return self._value

406

407

408class Gauge:

409 def __init__(self, name: str, description: str = ""):

410 self.name = name

411 self.description = description

412 self._value: float = 0.0

413

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

415 self._value = value

416

417 @property

418 def value(self) -> float:

419 return self._value

420

421

422class Histogram:

423 DEFAULT_BUCKETS = [1, 5, 10, 25, 50, 100, 250, 500, 1000, 5000, 10000]

424

425 def __init__(self, name: str, description: str = "", buckets: list[float] | None = None):

426 self.name = name

427 self.description = description

428 self.buckets: dict[float, int] = {}

429 for b in buckets or self.DEFAULT_BUCKETS:

430 self.buckets[b] = 0

431 self._count: int = 0

432 self._sum: float = 0.0

433

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

435 self._count += 1

436 self._sum += value

437 for boundary in sorted(self.buckets.keys()):

438 if value <= boundary:

439 self.buckets[boundary] += 1

440 break

441

442 @property

443 def count(self) -> int:

444 return self._count

445

446 @property

447 def sum(self) -> float:

448 return self._sum

449

450

451# ── Structured Logger ───────────────────────

452

453

454class LogLevel(StrEnum):

455 DEBUG = "debug"

456 INFO = "info"

457 WARNING = "warning"

458 ERROR = "error"

459 CRITICAL = "critical"

460

461

462class StructuredLogger:

463 """结构化 JSON 日志记录器，自动注入 trace_id 和 span_id。

464

465 Usage:

466 logger = StructuredLogger(tracer)

467 logger.info("Agent started", agent_name="ToolAgent", user_id="123")

468 """

469

470 def __init__(self, tracer: Tracer | None = None):

471 self._tracer = tracer

472 self._handlers: list[Callable] = []

473 self._level = LogLevel.INFO

474

475 def add_handler(self, handler: Callable[[dict], None]) -> None:

476 self._handlers.append(handler)

477

478 def set_level(self, level: LogLevel) -> None:

479 self._level = level

480

481 def log(self, level: LogLevel, message: str, **kwargs) -> None:

482 entry = {

483 "timestamp": time.time(),

484 "level": level.value,

485 "message": message,

486 **kwargs,

487 }

488

489 # Inject trace context

490 if self._tracer:

491 trace = self._tracer._active_trace

492 if trace:

493 entry["trace_id"] = trace.trace_id

494 if self._tracer._span_stack:

495 entry["span_id"] = self._tracer._span_stack[-1].span_id

496

497 for handler in self._handlers:

498 try:

499 handler(entry)

500 except Exception:

501 pass

502

503 def debug(self, message: str, **kwargs) -> None:

504 self.log(LogLevel.DEBUG, message, **kwargs)

505

506 def info(self, message: str, **kwargs) -> None:

507 self.log(LogLevel.INFO, message, **kwargs)

508

509 def warning(self, message: str, **kwargs) -> None:

510 self.log(LogLevel.WARNING, message, **kwargs)

511

512 def error(self, message: str, **kwargs) -> None:

513 self.log(LogLevel.ERROR, message, **kwargs)

514

515 def critical(self, message: str, **kwargs) -> None:

516 self.log(LogLevel.CRITICAL, message, **kwargs)

517

518

519# ── Dashboard Data ──────────────────────────

520

521

522class ObservabilityDashboard:

523 """可观测性数据聚合器 — 为 Grafana/自定义仪表盘提供数据。

524

525 Usage:

526 dashboard = ObservabilityDashboard(tracer, metrics)

527 summary = dashboard.get_summary()

528 """

529

530 def __init__(self, tracer: Tracer, metrics: MetricsRegistry):

531 self._tracer = tracer

532 self._metrics = metrics

533 self._trace_buffer: list[Trace] = []

534 self._max_buffer = 1000

535

536 def record_trace(self, trace: Trace) -> None:

537 self._trace_buffer.append(trace)

538 if len(self._trace_buffer) > self._max_buffer:

539 self._trace_buffer = self._trace_buffer[-self._max_buffer :]

540

541 def get_summary(self) -> dict:

542 """获取综合摘要。"""

543 traces = self._trace_buffer[-100:] # Last 100 traces

544

545 if not traces:

546 return {"message": "No traces recorded"}

547

548 durations = [t.total_duration_ms for t in traces]

549 costs = [t.total_cost_usd for t in traces]

550

551 durations.sort()

552

553 # Span kind distribution

554 kind_counts: dict[str, int] = defaultdict(int)

555 error_count = 0

556 for t in traces:

557 for s in t.spans:

558 kind_counts[s.kind.value] += 1

559 if s.status == SpanStatus.ERROR:

560 error_count += 1

561

562 return {

563 "trace_count": len(traces),

564 "total_traces": self._tracer._trace_count,

565 "error_rate": error_count / max(sum(kind_counts.values()), 1),

566 "duration": {

567 "p50_ms": durations[len(durations) // 2] if durations else 0,

568 "p95_ms": durations[int(len(durations) * 0.95)] if len(durations) > 1 else 0,

569 "p99_ms": durations[int(len(durations) * 0.99)] if len(durations) > 1 else 0,

570 "avg_ms": sum(durations) / len(durations) if durations else 0,

571 },

572 "cost": {

573 "total_usd": sum(costs),

574 "avg_per_trace_usd": sum(costs) / len(costs) if costs else 0,

575 },

576 "span_distribution": dict(kind_counts),

577 "metrics": self._metrics.collect(),

578 }

579

580 def get_latency_breakdown(self) -> list[dict]:

581 """按 pipeline 阶段拆分延迟。"""

582 breakdown: dict[str, list[float]] = defaultdict(list)

583

584 for trace in self._trace_buffer[-100:]:

585 for span in trace.spans:

586 breakdown[span.kind.value].append(span.duration_ms)

587

588 result = []

589 for kind, values in breakdown.items():

590 values.sort()

591 n = len(values)

592 result.append(

593 {

594 "stage": kind,

595 "count": n,

596 "avg_ms": sum(values) / n if n else 0,

597 "p50_ms": values[n // 2] if n else 0,

598 "p95_ms": values[int(n * 0.95)] if n > 1 else (values[0] if values else 0),

599 }

600 )

601

602 return result

603

604

605# ── Quick Start ─────────────────────────────

606

607

608def create_observability_stack(service_name: str = "agentos"):

609 """一键创建可观测性栈。"""

610 tracer = Tracer(service_name=service_name)

611 metrics = MetricsRegistry()

612 logger = StructuredLogger(tracer)

613 dashboard = ObservabilityDashboard(tracer, metrics)

614

615 # Register auto-export

616 tracer.add_exporter(dashboard.record_trace)

617

618 # Register default metrics

619 metrics.counter("agent_invocations_total", "Total agent invocations")

620 metrics.histogram("agent_latency_ms", "Agent end-to-end latency")

621 metrics.histogram("llm_latency_ms", "LLM call latency")

622 metrics.counter("tool_calls_total", "Total tool calls")

623 metrics.counter("tool_errors_total", "Total tool errors")

624 metrics.gauge("active_agents", "Currently active agents")

625

626 return tracer, metrics, logger, dashboard

627

628

629# ── Missing compat classes ───────────────────

630

631

632class MetricsCollector(MetricsRegistry):

633 """Alias for MetricsRegistry — required by agentos/__init__.py."""

634

635

636

637class NoopTracer(Tracer):

638 """No-op tracer — compatible replacement."""

639

640 def __init__(self, *args, **kwargs):

641 super().__init__(*args, **kwargs)

642

643 def start_span(self, *args, **kwargs):

644 return Span(name="noop")

645

646 def mark_complete(self, span_id: str = ""):

647 pass

648

649

650@dataclass

651class CostAnalytics:

652 """LLM 成本分析。"""

653

654 total_cost: float = 0.0

655 total_tokens: int = 0

656 prompt_tokens: int = 0

657 completion_tokens: int = 0

658 cost_per_call: list[float] = field(default_factory=list)

659

660 def record_call(self, prompt_tokens: int, completion_tokens: int, cost: float):

661 self.prompt_tokens += prompt_tokens

662 self.completion_tokens += completion_tokens

663 self.total_cost += cost

664 self.cost_per_call.append(cost)

665

666 def summary(self) -> dict:

667 return {

668 "total_cost": self.total_cost,

669 "total_tokens": self.prompt_tokens + self.completion_tokens,

670 "avg_cost_per_call": (

671 sum(self.cost_per_call) / len(self.cost_per_call) if self.cost_per_call else 0

672 ),

673 }

674

675

676@dataclass

677class BudgetAlert:

678 """成本预算告警。"""

679

680 budget_limit: float = 0.0

681 current_spend: float = 0.0

682 threshold_pct: float = 80.0

683

684 def check(self) -> str | None:

685 if self.budget_limit <= 0:

686 return None

687 pct = (self.current_spend / self.budget_limit) * 100

688 if pct >= self.threshold_pct:

689 return f"Budget alert: {pct:.1f}% of ${self.budget_limit:.2f} spent"

690 return None

691

692

693# ── Re-export tracing module for convenience ───────────────────────────

694

Coverage for agentos/observability/init.py: 42%

352 statements