computer-agent-py 0.1.0__py3-none-any.whl

computeragent/telemetry/sinks/otel.py

@@ -0,0 +1,375 @@
+"""``OtelSink`` — emit GenAI-semconv-compliant OpenTelemetry spans + metrics.
+
+Span tree matches the TS ``@computeragent/observability`` package so polyglot
+deployments produce one coherent trace tree in the same backend::
+
+    invoke_agent <agent_name>   [gen_ai.conversation.id = session_id,
+                                 gen_ai.system = "claude-agent-sdk",
+                                 gen_ai.agent.name = ...]
+    └── chat <model>            [gen_ai.request.model = ...]
+        └── execute_tool <name> [gen_ai.tool.name = ...,
+                                 gen_ai.tool.call.id = ...]
+
+Metrics::
+
+    gen_ai.client.token.usage           (histogram, unit "{token}")
+        dims: gen_ai.operation.name, gen_ai.request.model, gen_ai.token.type
+    gen_ai.client.operation.duration    (histogram, unit "s")
+    computeragent.usage.cost_usd        (histogram, unit "{usd}") — custom
+"""
+
+from __future__ import annotations
+
+import contextlib
+import logging
+from typing import TYPE_CHECKING, Any
+
+try:
+    from opentelemetry import metrics, trace  # noqa: F401 - probe imports only
+    from opentelemetry.exporter.otlp.proto.http.metric_exporter import (
+        OTLPMetricExporter as HttpMetricExporter,
+    )
+    from opentelemetry.exporter.otlp.proto.http.trace_exporter import (
+        OTLPSpanExporter as HttpSpanExporter,
+    )
+    from opentelemetry.sdk.metrics import MeterProvider
+    from opentelemetry.sdk.metrics.export import (
+        ConsoleMetricExporter,
+        PeriodicExportingMetricReader,
+    )
+    from opentelemetry.sdk.resources import Resource
+    from opentelemetry.sdk.trace import TracerProvider
+    from opentelemetry.sdk.trace.export import (
+        BatchSpanProcessor,
+        ConsoleSpanExporter,
+        SimpleSpanProcessor,
+    )
+except ImportError as _exc:  # pragma: no cover - guarded by lazy export
+    raise ImportError(
+        "OtelSink requires the [otel] extra. Install with: pip install 'computer-agent-py[otel]'"
+    ) from _exc
+
+if TYPE_CHECKING:
+    from opentelemetry.trace import Span
+
+    from ..event import TelemetryEvent
+
+logger = logging.getLogger("computeragent.telemetry")
+
+
+# GenAI semconv attribute keys — pinned literals to guarantee parity with the
+# TS @computeragent/observability emission. The semconv Python package's
+# constants for these keys move between alpha versions; literals are safer.
+GEN_AI_SYSTEM = "gen_ai.system"
+GEN_AI_AGENT_NAME = "gen_ai.agent.name"
+GEN_AI_CONVERSATION_ID = "gen_ai.conversation.id"
+GEN_AI_REQUEST_MODEL = "gen_ai.request.model"
+GEN_AI_OPERATION_NAME = "gen_ai.operation.name"
+GEN_AI_TOOL_NAME = "gen_ai.tool.name"
+GEN_AI_TOOL_CALL_ID = "gen_ai.tool.call.id"
+GEN_AI_TOKEN_TYPE = "gen_ai.token.type"
+GEN_AI_USAGE_INPUT_TOKENS = "gen_ai.usage.input_tokens"
+GEN_AI_USAGE_OUTPUT_TOKENS = "gen_ai.usage.output_tokens"
+GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS = "gen_ai.usage.cache_read_input_tokens"
+GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS = "gen_ai.usage.cache_creation_input_tokens"
+
+METRIC_TOKEN_USAGE = "gen_ai.client.token.usage"
+METRIC_OPERATION_DURATION = "gen_ai.client.operation.duration"
+METRIC_COST_USD = "computeragent.usage.cost_usd"
+
+
+class OtelSink:
+    """Translate :class:`TelemetryEvent` into OTel spans + metric measurements.
+
+    Configures its own ``TracerProvider`` and ``MeterProvider`` on first use
+    unless ``tracer_provider`` / ``meter_provider`` is passed in (tests use
+    this hook to attach in-memory exporters).
+    """
+
+    def __init__(
+        self,
+        *,
+        service_name: str | None = None,
+        exporter: str | None = None,
+        endpoint: str | None = None,
+        headers: dict[str, str] | None = None,
+        tracer_provider: Any | None = None,
+        meter_provider: Any | None = None,
+    ) -> None:
+        # Read effective config from the module-level resolver if no overrides.
+        from ..config import get_config
+
+        cfg = get_config()
+        self._service_name = service_name or cfg.get("service_name") or "computeragent"
+        self._exporter = (
+            exporter or cfg.get("exporter") or _infer_exporter(endpoint or cfg.get("endpoint"))
+        )
+        self._endpoint = endpoint or cfg.get("endpoint")
+        self._headers = headers or cfg.get("headers") or {}
+
+        self._tracer_provider = tracer_provider or _build_tracer_provider(
+            service_name=self._service_name,
+            exporter=self._exporter,
+            endpoint=self._endpoint,
+            headers=self._headers,
+        )
+        self._meter_provider = meter_provider or _build_meter_provider(
+            service_name=self._service_name,
+            exporter=self._exporter,
+            endpoint=self._endpoint,
+            headers=self._headers,
+        )
+        self._tracer = self._tracer_provider.get_tracer("computeragent")
+        self._meter = self._meter_provider.get_meter("computeragent")
+
+        self._token_usage_hist = self._meter.create_histogram(
+            METRIC_TOKEN_USAGE, unit="{token}", description="GenAI token usage"
+        )
+        self._duration_hist = self._meter.create_histogram(
+            METRIC_OPERATION_DURATION, unit="s", description="GenAI operation duration"
+        )
+        self._cost_hist = self._meter.create_histogram(
+            METRIC_COST_USD, unit="{usd}", description="Inference cost in USD"
+        )
+
+        # session_id → (invoke_agent span, chat span, active execute_tool span)
+        # Spans are held open across events for the same session/turn.
+        self._invoke_spans: dict[str, Span] = {}
+        self._chat_spans: dict[str, Span] = {}
+        self._tool_spans: dict[str, dict[str, Span]] = {}  # session → tool_use_id → span
+
+    # ── lifecycle ----------------------------------------------------------
+
+    async def emit(self, event: TelemetryEvent) -> None:
+        handler = _DISPATCH.get(event.kind)
+        if handler is None:
+            return
+        handler(self, event)
+
+    async def flush(self, timeout: float = 5.0) -> None:
+        # OTel SDK's flush takes milliseconds. Convert + fan out.
+        ms = int(timeout * 1000)
+        try:
+            self._tracer_provider.force_flush(timeout_millis=ms)
+        except Exception:  # noqa: BLE001
+            logger.debug("tracer_provider.force_flush failed", exc_info=True)
+        try:
+            self._meter_provider.force_flush(timeout_millis=ms)
+        except Exception:  # noqa: BLE001
+            logger.debug("meter_provider.force_flush failed", exc_info=True)
+
+    async def aclose(self) -> None:
+        await self.flush()
+        try:
+            self._tracer_provider.shutdown()
+        except Exception:  # noqa: BLE001
+            logger.debug("tracer_provider.shutdown failed", exc_info=True)
+        try:
+            self._meter_provider.shutdown()
+        except Exception:  # noqa: BLE001
+            logger.debug("meter_provider.shutdown failed", exc_info=True)
+
+    # ── handlers (one per event kind) --------------------------------------
+
+    def _on_session_started(self, event: TelemetryEvent) -> None:
+        span = self._tracer.start_span(name=f"invoke_agent {event.agent_name or 'agent'}")
+        span.set_attribute(GEN_AI_SYSTEM, "claude-agent-sdk")
+        if event.agent_name:
+            span.set_attribute(GEN_AI_AGENT_NAME, event.agent_name)
+        span.set_attribute(GEN_AI_CONVERSATION_ID, event.session_id)
+        span.set_attribute(GEN_AI_OPERATION_NAME, "invoke_agent")
+        if event.payload.get("model"):
+            span.set_attribute(GEN_AI_REQUEST_MODEL, str(event.payload["model"]))
+        self._invoke_spans[event.session_id] = span
+
+    def _on_assistant_message(self, event: TelemetryEvent) -> None:
+        # Open or reuse a chat span; spec models chat as one span per LLM call,
+        # but the upstream stream gives us one AssistantMessage per LLM call,
+        # so open/close per event.
+        chat = self._tracer.start_span(name=f"chat {event.payload.get('model') or 'unknown'}")
+        chat.set_attribute(GEN_AI_OPERATION_NAME, "chat")
+        if event.payload.get("model"):
+            chat.set_attribute(GEN_AI_REQUEST_MODEL, str(event.payload["model"]))
+        chat.set_attribute(GEN_AI_CONVERSATION_ID, event.session_id)
+        chat.end()
+
+    def _on_tool_use(self, event: TelemetryEvent) -> None:
+        tool_name = str(event.payload.get("tool_name") or "tool")
+        span = self._tracer.start_span(name=f"execute_tool {tool_name}")
+        span.set_attribute(GEN_AI_OPERATION_NAME, "execute_tool")
+        span.set_attribute(GEN_AI_TOOL_NAME, tool_name)
+        call_id = str(event.payload.get("tool_use_id") or "")
+        if call_id:
+            span.set_attribute(GEN_AI_TOOL_CALL_ID, call_id)
+        span.set_attribute(GEN_AI_CONVERSATION_ID, event.session_id)
+        self._tool_spans.setdefault(event.session_id, {})[call_id] = span
+
+    def _on_tool_result(self, event: TelemetryEvent) -> None:
+        call_id = str(event.payload.get("tool_use_id") or "")
+        span = self._tool_spans.get(event.session_id, {}).pop(call_id, None)
+        if span is None:
+            return
+        if event.payload.get("is_error"):
+            span.set_attribute("error.type", "tool_error")
+        span.end()
+
+    def _on_usage_snapshot(self, event: TelemetryEvent) -> None:
+        usage = event.payload.get("usage") or {}
+        model = event.payload.get("model")
+        common_dims = {
+            GEN_AI_OPERATION_NAME: "chat",
+            **({GEN_AI_REQUEST_MODEL: str(model)} if model else {}),
+        }
+        for token_type, key in (
+            ("input", "input_tokens"),
+            ("output", "output_tokens"),
+            ("cache_read", "cache_read_input_tokens"),
+            ("cache_creation", "cache_creation_input_tokens"),
+        ):
+            raw = usage.get(key)
+            if raw is None:
+                continue
+            try:
+                v = int(raw)
+            except (TypeError, ValueError):
+                continue
+            self._token_usage_hist.record(
+                v, attributes={**common_dims, GEN_AI_TOKEN_TYPE: token_type}
+            )
+
+        cost = event.payload.get("total_cost_usd")
+        if cost is not None:
+            with contextlib.suppress(TypeError, ValueError):
+                self._cost_hist.record(float(cost), attributes=common_dims)
+
+    def _on_session_ended(self, event: TelemetryEvent) -> None:
+        # Close any tool spans still open for this session (defensive — should
+        # already be closed by tool_result events).
+        for span in self._tool_spans.pop(event.session_id, {}).values():
+            try:
+                span.end()
+            except Exception:  # noqa: BLE001
+                logger.debug("stray tool span end failed", exc_info=True)
+
+        invoke = self._invoke_spans.pop(event.session_id, None)
+        if invoke is None:
+            return
+        if event.payload.get("is_error"):
+            invoke.set_attribute(
+                "error.type",
+                str(event.payload.get("error_type") or event.payload.get("subtype") or "error"),
+            )
+        usage = event.payload.get("usage") or {}
+        for key, attr in (
+            ("input_tokens", GEN_AI_USAGE_INPUT_TOKENS),
+            ("output_tokens", GEN_AI_USAGE_OUTPUT_TOKENS),
+            ("cache_read_input_tokens", GEN_AI_USAGE_CACHE_READ_INPUT_TOKENS),
+            ("cache_creation_input_tokens", GEN_AI_USAGE_CACHE_CREATION_INPUT_TOKENS),
+        ):
+            v = usage.get(key)
+            if isinstance(v, int):
+                invoke.set_attribute(attr, v)
+        invoke.end()
+        duration_s = float(event.payload.get("duration_ms", 0.0) or 0.0) / 1000.0
+        if duration_s > 0:
+            self._duration_hist.record(
+                duration_s, attributes={GEN_AI_OPERATION_NAME: "invoke_agent"}
+            )
+
+    def _on_user_message(self, _event: TelemetryEvent) -> None:
+        # No span; user messages are inputs, recorded as conversation events
+        # on the active invoke_agent span in a future iteration.
+        return
+
+    def _on_system_message(self, _event: TelemetryEvent) -> None:
+        return
+
+    def _on_policy_decision(self, event: TelemetryEvent) -> None:
+        # Annotate the active execute_tool span for this session, if any.
+        # Otherwise drop a span event onto the invoke_agent span so the
+        # decision is still visible in the trace.
+        attrs: dict[str, Any] = {
+            "policy.engine": event.payload.get("engine", ""),
+            "policy.decision": event.payload.get("decision", ""),
+            "policy.reason": event.payload.get("reason", ""),
+            "policy.latency_ms": float(event.payload.get("latency_ms", 0.0) or 0.0),
+            "policy.tool_name": event.payload.get("tool_name", ""),
+        }
+        tool_name = str(event.payload.get("tool_name") or "")
+        tool_spans = self._tool_spans.get(event.session_id, {})
+        # Tool spans are keyed by tool_use_id, which we don't have here. Pick
+        # the most-recently-opened span with a matching tool name.
+        target_span = None
+        for span in tool_spans.values():
+            if getattr(span, "name", "").endswith(tool_name) and tool_name:
+                target_span = span
+                break
+        if target_span is None:
+            target_span = self._invoke_spans.get(event.session_id)
+        if target_span is not None:
+            try:
+                target_span.add_event("policy_decision", attributes=attrs)
+            except Exception:  # noqa: BLE001
+                logger.debug("policy_decision span event failed", exc_info=True)
+            return
+        # No active span — emit a standalone span so the decision is captured.
+        try:
+            span = self._tracer.start_span(name=f"policy_decision {tool_name or 'tool'}")
+            for k, v in attrs.items():
+                span.set_attribute(k, v)
+            span.set_attribute(GEN_AI_CONVERSATION_ID, event.session_id)
+            span.end()
+        except Exception:  # noqa: BLE001
+            logger.debug("policy_decision standalone span failed", exc_info=True)
+
+
+# Map kind → instance method. Resolved at module import; no per-event isinstance.
+_DISPATCH: dict[str, Any] = {
+    "session_started": OtelSink._on_session_started,
+    "assistant_message": OtelSink._on_assistant_message,
+    "tool_use": OtelSink._on_tool_use,
+    "tool_result": OtelSink._on_tool_result,
+    "usage_snapshot": OtelSink._on_usage_snapshot,
+    "session_ended": OtelSink._on_session_ended,
+    "user_message": OtelSink._on_user_message,
+    "system_message": OtelSink._on_system_message,
+    "policy_decision": OtelSink._on_policy_decision,
+}
+
+
+def _infer_exporter(endpoint: str | None) -> str:
+    return "otlp-http" if endpoint else "console"
+
+
+def _build_tracer_provider(
+    *,
+    service_name: str,
+    exporter: str,
+    endpoint: str | None,
+    headers: dict[str, str],
+) -> TracerProvider:
+    resource = Resource.create({"service.name": service_name})
+    provider = TracerProvider(resource=resource)
+    if exporter == "otlp-http":
+        span_exporter = HttpSpanExporter(endpoint=endpoint, headers=headers or None)
+        provider.add_span_processor(BatchSpanProcessor(span_exporter))
+    else:
+        provider.add_span_processor(SimpleSpanProcessor(ConsoleSpanExporter()))
+    return provider
+
+
+def _build_meter_provider(
+    *,
+    service_name: str,
+    exporter: str,
+    endpoint: str | None,
+    headers: dict[str, str],
+) -> MeterProvider:
+    resource = Resource.create({"service.name": service_name})
+    metric_exporter: Any
+    if exporter == "otlp-http":
+        metric_exporter = HttpMetricExporter(endpoint=endpoint, headers=headers or None)
+    else:
+        metric_exporter = ConsoleMetricExporter()
+    reader = PeriodicExportingMetricReader(metric_exporter, export_interval_millis=5000)
+    return MeterProvider(resource=resource, metric_readers=[reader])

computeragent/types.py

@@ -0,0 +1,13 @@
+"""Re-export of ``claude_agent_sdk.types``.
+
+Every type that ``claude_agent_sdk.types`` exposes is re-exported here with
+identity preserved. This is what makes ``isinstance(msg, ResultMessage)`` work
+whether the user imported ``ResultMessage`` from ``claude_agent_sdk.types`` or
+from ``computeragent.types``.
+
+Switching the import line is the only change a drop-in user needs.
+"""
+
+from __future__ import annotations
+
+from claude_agent_sdk.types import *  # noqa: F401, F403