GenAI Utils | Agent metrics and events

Author: etserend

util/opentelemetry-util-genai/src/opentelemetry/util/genai/handler.py

@@ -60,9 +60,10 @@
 
 from __future__ import annotations
 
+import logging
 import timeit
 from contextlib import contextmanager
-from typing import Iterator
+from typing import Callable, Iterator, TypeVar
 
 from opentelemetry import context as otel_context
 from opentelemetry._logs import (
@@ -80,13 +81,60 @@
 )
 from opentelemetry.util.genai.metrics import InvocationMetricsRecorder
 from opentelemetry.util.genai.span_utils import (
+    _apply_agent_finish_attributes,
+    _apply_creation_finish_attributes,
     _apply_error_attributes,
     _apply_llm_finish_attributes,
+    _maybe_emit_agent_event,
     _maybe_emit_llm_event,
 )
-from opentelemetry.util.genai.types import Error, LLMInvocation
+from opentelemetry.util.genai.types import (
+    AgentCreation,
+    AgentInvocation,
+    Error,
+    GenAIInvocation,
+    LLMInvocation,
+)
 from opentelemetry.util.genai.version import __version__
 
+_logger = logging.getLogger(__name__)
+
+_T = TypeVar("_T", bound=GenAIInvocation)
+
+
+@contextmanager
+def _lifecycle_context(
+    invocation: _T,
+    start: Callable[[_T], _T],
+    stop: Callable[[_T], _T],
+    fail: Callable[[_T, Error], _T],
+    label: str,
+) -> Iterator[_T]:
+    """Shared lifecycle context manager for GenAI invocations.
+
+    Wraps start/stop/fail calls with error handling so SDK-internal
+    errors never propagate to the caller.
+    """
+    try:
+        start(invocation)
+    except Exception:
+        _logger.warning("Failed to start %s span", label, exc_info=True)
+    try:
+        yield invocation
+    except Exception as exc:
+        try:
+            fail(invocation, Error(message=str(exc), type=type(exc)))
+        except Exception:
+            _logger.warning(
+                "Failed to record %s failure", label, exc_info=True
+            )
+        raise
+    else:
+        try:
+            stop(invocation)
+        except Exception:
+            _logger.warning("Failed to stop %s span", label, exc_info=True)
+
 
 class TelemetryHandler:
     """
@@ -156,13 +204,13 @@ def stop_llm(self, invocation: LLMInvocation) -> LLMInvocation:  # pylint: disab
             # TODO: Provide feedback that this invocation was not started
             return invocation
 
-        span = invocation.span
-        _apply_llm_finish_attributes(span, invocation)
-        self._record_llm_metrics(invocation, span)
-        _maybe_emit_llm_event(self._logger, span, invocation)
-        # Detach context and end span
-        otel_context.detach(invocation.context_token)
-        span.end()
+        try:
+            _apply_llm_finish_attributes(invocation.span, invocation)
+            self._record_llm_metrics(invocation, invocation.span)
+            _maybe_emit_llm_event(self._logger, invocation.span, invocation)
+        finally:
+            otel_context.detach(invocation.context_token)
+            invocation.span.end()
         return invocation
 
     def fail_llm(  # pylint: disable=no-self-use
@@ -173,15 +221,19 @@ def fail_llm(  # pylint: disable=no-self-use
             # TODO: Provide feedback that this invocation was not started
             return invocation
 
-        span = invocation.span
-        _apply_llm_finish_attributes(invocation.span, invocation)
-        _apply_error_attributes(invocation.span, error)
-        error_type = getattr(error.type, "__qualname__", None)
-        self._record_llm_metrics(invocation, span, error_type=error_type)
-        _maybe_emit_llm_event(self._logger, span, invocation, error)
-        # Detach context and end span
-        otel_context.detach(invocation.context_token)
-        span.end()
+        try:
+            _apply_llm_finish_attributes(invocation.span, invocation)
+            _apply_error_attributes(invocation.span, error)
+            error_type = getattr(error.type, "__qualname__", None)
+            self._record_llm_metrics(
+                invocation, invocation.span, error_type=error_type
+            )
+            _maybe_emit_llm_event(
+                self._logger, invocation.span, invocation, error
+            )
+        finally:
+            otel_context.detach(invocation.context_token)
+            invocation.span.end()
         return invocation
 
     @contextmanager
@@ -200,13 +252,173 @@ def llm(
             invocation = LLMInvocation(
                 request_model="",
             )
-        self.start_llm(invocation)
+        with _lifecycle_context(
+            invocation, self.start_llm, self.stop_llm, self.fail_llm, "llm"
+        ) as inv:
+            yield inv
+
+    # ---- Agent invocation lifecycle ----
+
+    def start_agent(
+        self,
+        invocation: AgentInvocation,
+    ) -> AgentInvocation:
+        """Start an agent invocation and create a pending span entry."""
+        span_name = (
+            f"{invocation.operation_name} {invocation.agent_name}".strip()
+        )
+        kind = SpanKind.CLIENT if invocation.is_remote else SpanKind.INTERNAL
+        span = self._tracer.start_span(
+            name=span_name,
+            kind=kind,
+        )
+        invocation.monotonic_start_s = timeit.default_timer()
+        invocation.span = span
+        invocation.context_token = otel_context.attach(
+            set_span_in_context(span)
+        )
+        return invocation
+
+    def _record_agent_metrics(
+        self,
+        invocation: AgentInvocation,
+        span: Span | None = None,
+        *,
+        error_type: str | None = None,
+    ) -> None:
+        if self._metrics_recorder is None or span is None:
+            return
+        self._metrics_recorder.record_agent(
+            span,
+            invocation,
+            error_type=error_type,
+        )
+
+    def stop_agent(self, invocation: AgentInvocation) -> AgentInvocation:
+        """Finalize an agent invocation successfully and end its span."""
+        if invocation.context_token is None or invocation.span is None:
+            return invocation
+
+        try:
+            _apply_agent_finish_attributes(invocation.span, invocation)
+            self._record_agent_metrics(invocation, invocation.span)
+            _maybe_emit_agent_event(self._logger, invocation.span, invocation)
+        finally:
+            otel_context.detach(invocation.context_token)
+            invocation.span.end()
+        return invocation
+
+    def fail_agent(
+        self, invocation: AgentInvocation, error: Error
+    ) -> AgentInvocation:
+        """Fail an agent invocation and end its span with error status."""
+        if invocation.context_token is None or invocation.span is None:
+            return invocation
+
+        try:
+            _apply_agent_finish_attributes(invocation.span, invocation)
+            _apply_error_attributes(invocation.span, error)
+            error_type = getattr(error.type, "__qualname__", None)
+            self._record_agent_metrics(
+                invocation, invocation.span, error_type=error_type
+            )
+            _maybe_emit_agent_event(
+                self._logger, invocation.span, invocation, error
+            )
+        finally:
+            otel_context.detach(invocation.context_token)
+            invocation.span.end()
+        return invocation
+
+    @contextmanager
+    def agent(
+        self, invocation: AgentInvocation | None = None
+    ) -> Iterator[AgentInvocation]:
+        """Context manager for agent invocations.
+
+        Only set data attributes on the invocation object, do not modify the span or context.
+
+        Starts the span on entry. On normal exit, finalizes the invocation and ends the span.
+        If an exception occurs inside the context, marks the span as error, ends it, and
+        re-raises the original exception.
+        """
+        if invocation is None:
+            invocation = AgentInvocation()
+        with _lifecycle_context(
+            invocation,
+            self.start_agent,
+            self.stop_agent,
+            self.fail_agent,
+            "agent",
+        ) as inv:
+            yield inv
+
+    # ---- Agent creation lifecycle ----
+
+    def start_create_agent(
+        self,
+        creation: AgentCreation,
+    ) -> AgentCreation:
+        """Start an agent creation and create a pending span entry."""
+        span_name = f"{creation.operation_name} {creation.agent_name}".strip()
+        span = self._tracer.start_span(
+            name=span_name,
+            kind=SpanKind.CLIENT,
+        )
+        creation.monotonic_start_s = timeit.default_timer()
+        creation.span = span
+        creation.context_token = otel_context.attach(set_span_in_context(span))
+        return creation
+
+    def stop_create_agent(self, creation: AgentCreation) -> AgentCreation:  # pylint: disable=no-self-use
+        """Finalize an agent creation successfully and end its span."""
+        if creation.context_token is None or creation.span is None:
+            return creation
+
         try:
-            yield invocation
-        except Exception as exc:
-            self.fail_llm(invocation, Error(message=str(exc), type=type(exc)))
-            raise
-        self.stop_llm(invocation)
+            _apply_creation_finish_attributes(creation.span, creation)
+        finally:
+            otel_context.detach(creation.context_token)
+            creation.span.end()
+        return creation
+
+    def fail_create_agent(  # pylint: disable=no-self-use
+        self, creation: AgentCreation, error: Error
+    ) -> AgentCreation:
+        """Fail an agent creation and end its span with error status."""
+        if creation.context_token is None or creation.span is None:
+            return creation
+
+        try:
+            _apply_creation_finish_attributes(creation.span, creation)
+            _apply_error_attributes(creation.span, error)
+        finally:
+            otel_context.detach(creation.context_token)
+            creation.span.end()
+        return creation
+
+    @contextmanager
+    def create_agent(
+        self, creation: AgentCreation | None = None
+    ) -> Iterator[AgentCreation]:
+        """Context manager for agent creation.
+
+        Only set data attributes on the creation object, do not modify the span or context.
+
+        Starts the span on entry. On normal exit, finalizes the creation and ends the span.
+        If an exception occurs inside the context, marks the span as error, ends it, and
+        re-raises the original exception.
+        """
+        if creation is None:
+            creation = AgentCreation()
+        with _lifecycle_context(
+            creation,
+            self.start_create_agent,
+            self.stop_create_agent,
+            self.fail_create_agent,
+            "create_agent",
+        ) as c:
+            yield c
 
 
 def get_telemetry_handler(

util/opentelemetry-util-genai/src/opentelemetry/util/genai/metrics.py

@@ -18,7 +18,7 @@
     create_duration_histogram,
     create_token_histogram,
 )
-from opentelemetry.util.genai.types import LLMInvocation
+from opentelemetry.util.genai.types import AgentInvocation, LLMInvocation
 from opentelemetry.util.types import AttributeValue
 
 
@@ -105,5 +105,78 @@ def record(
                 context=span_context,
             )
 
+    def record_agent(
+        self,
+        span: Optional[Span],
+        invocation: AgentInvocation,
+        *,
+        error_type: Optional[str] = None,
+    ) -> None:
+        """Record duration and token metrics for an agent invocation."""
+
+        if span is None:
+            return
+
+        token_counts: list[tuple[int, str]] = []
+        if invocation.input_tokens is not None:
+            token_counts.append(
+                (
+                    invocation.input_tokens,
+                    GenAI.GenAiTokenTypeValues.INPUT.value,
+                )
+            )
+        if invocation.output_tokens is not None:
+            token_counts.append(
+                (
+                    invocation.output_tokens,
+                    GenAI.GenAiTokenTypeValues.OUTPUT.value,
+                )
+            )
+
+        attributes: Dict[str, AttributeValue] = {
+            GenAI.GEN_AI_OPERATION_NAME: invocation.operation_name
+        }
+        if invocation.request_model:
+            attributes[GenAI.GEN_AI_REQUEST_MODEL] = invocation.request_model
+        if invocation.provider:
+            attributes[GenAI.GEN_AI_PROVIDER_NAME] = invocation.provider
+        if invocation.response_model_name:
+            attributes[GenAI.GEN_AI_RESPONSE_MODEL] = (
+                invocation.response_model_name
+            )
+        if invocation.server_address:
+            attributes[server_attributes.SERVER_ADDRESS] = (
+                invocation.server_address
+            )
+        if invocation.server_port is not None:
+            attributes[server_attributes.SERVER_PORT] = invocation.server_port
+        if invocation.metric_attributes:
+            attributes.update(invocation.metric_attributes)
+
+        duration_seconds: Optional[float] = None
+        if invocation.monotonic_start_s is not None:
+            duration_seconds = max(
+                timeit.default_timer() - invocation.monotonic_start_s,
+                0.0,
+            )
+
+        span_context = set_span_in_context(span)
+        if error_type:
+            attributes[error_attributes.ERROR_TYPE] = error_type
+
+        if duration_seconds is not None:
+            self._duration_histogram.record(
+                duration_seconds,
+                attributes=attributes,
+                context=span_context,
+            )
+
+        for token_count, token_type in token_counts:
+            self._token_histogram.record(
+                token_count,
+                attributes=attributes | {GenAI.GEN_AI_TOKEN_TYPE: token_type},
+                context=span_context,
+            )
+
 
 __all__ = ["InvocationMetricsRecorder"]