unify handler start/stop/fail

Author: etserend

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

@@ -62,7 +62,7 @@
 
 import timeit
 from contextlib import contextmanager
-from typing import Iterator
+from typing import Iterator, TypeVar
 
 from opentelemetry import context as otel_context
 from opentelemetry._logs import (
@@ -78,6 +78,7 @@
     get_tracer,
     set_span_in_context,
 )
+from opentelemetry.trace.status import Status, StatusCode
 from opentelemetry.util.genai.metrics import InvocationMetricsRecorder
 from opentelemetry.util.genai.span_utils import (
     _apply_agent_finish_attributes,
@@ -90,10 +91,13 @@
     AgentCreation,
     AgentInvocation,
     Error,
+    GenAIInvocation,
     LLMInvocation,
 )
 from opentelemetry.util.genai.version import __version__
 
+_T = TypeVar("_T", bound=GenAIInvocation)
+
 
 class TelemetryHandler:
     """
@@ -138,14 +142,10 @@ def _record_llm_metrics(
             error_type=error_type,
         )
 
-    def start_llm(
-        self,
-        invocation: LLMInvocation,
-    ) -> LLMInvocation:
-        """Start an LLM invocation and create a pending span entry."""
-        # Create a span and attach it as current; keep the token to detach later
+    def _start(self, invocation: _T) -> _T:
+        """Start a GenAI invocation and create a pending span entry."""
         span = self._tracer.start_span(
-            name=f"{invocation.operation_name} {invocation.request_model}",
+            name=invocation.operation_name or "",
             kind=SpanKind.CLIENT,
         )
         # Record a monotonic start timestamp (seconds) for duration
@@ -157,40 +157,106 @@ def start_llm(
         )
         return invocation
 
-    def stop_llm(self, invocation: LLMInvocation) -> LLMInvocation:  # pylint: disable=no-self-use
-        """Finalize an LLM invocation successfully and end its span."""
+    def _stop(self, invocation: _T) -> _T:
+        """Finalize a GenAI invocation successfully and end its span."""
         if invocation.context_token is None or invocation.span is None:
             # 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:
+            if isinstance(invocation, LLMInvocation):
+                _apply_llm_finish_attributes(span, invocation)
+                self._record_llm_metrics(invocation, span)
+                _maybe_emit_llm_event(self._logger, span, invocation)
+            elif isinstance(invocation, AgentInvocation):
+                _apply_agent_finish_attributes(span, invocation)
+            elif isinstance(invocation, AgentCreation):
+                _apply_creation_finish_attributes(span, invocation)
+            else:
+                span.set_status(
+                    Status(
+                        StatusCode.ERROR,
+                        f"Unsupported invocation type: {type(invocation)!r}",
+                    )
+                )
+                raise TypeError(
+                    f"Unsupported invocation type: {type(invocation)!r}"
+                )
+        finally:
+            otel_context.detach(invocation.context_token)
+            span.end()
         return invocation
 
-    def fail_llm(  # pylint: disable=no-self-use
-        self, invocation: LLMInvocation, error: Error
-    ) -> LLMInvocation:
-        """Fail an LLM invocation and end its span with error status."""
+    def _fail(self, invocation: _T, error: Error) -> _T:
+        """Fail a GenAI invocation and end its span with error status."""
         if invocation.context_token is None or invocation.span is None:
             # 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:
+            if isinstance(invocation, LLMInvocation):
+                _apply_llm_finish_attributes(span, invocation)
+                _apply_error_attributes(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)
+            elif isinstance(invocation, AgentInvocation):
+                _apply_agent_finish_attributes(span, invocation)
+                _apply_error_attributes(span, error)
+            elif isinstance(invocation, AgentCreation):
+                _apply_creation_finish_attributes(span, invocation)
+                _apply_error_attributes(span, error)
+            else:
+                span.set_status(
+                    Status(
+                        StatusCode.ERROR,
+                        f"Unsupported invocation type: {type(invocation)!r}",
+                    )
+                )
+                raise TypeError(
+                    f"Unsupported invocation type: {type(invocation)!r}"
+                )
+        finally:
+            otel_context.detach(invocation.context_token)
+            span.end()
         return invocation
 
+    def start(
+        self,
+        invocation: _T,
+    ) -> _T:
+        """Start a GenAI invocation and create a pending span entry."""
+        return self._start(invocation)
+
+    def stop(self, invocation: _T) -> _T:
+        """Finalize a GenAI invocation successfully and end its span."""
+        return self._stop(invocation)
+
+    def fail(self, invocation: _T, error: Error) -> _T:
+        """Fail a GenAI invocation and end its span with error status."""
+        return self._fail(invocation, error)
+
+    def start_llm(
+        self,
+        invocation: LLMInvocation,
+    ) -> LLMInvocation:
+        """Start an LLM invocation and create a pending span entry."""
+        return self.start(invocation)
+
+    def stop_llm(self, invocation: LLMInvocation) -> LLMInvocation:
+        """Finalize an LLM invocation successfully and end its span."""
+        return self.stop(invocation)
+
+    def fail_llm(
+        self, invocation: LLMInvocation, error: Error
+    ) -> LLMInvocation:
+        """Fail an LLM invocation and end its span with error status."""
+        return self.fail(invocation, error)
+
     @contextmanager
     def llm(
         self, invocation: LLMInvocation | None = None
@@ -222,43 +288,17 @@ def start_agent(
         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
+        return self.start(invocation)
 
-    def stop_agent(self, invocation: AgentInvocation) -> AgentInvocation:  # pylint: disable=no-self-use
+    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
-
-        span = invocation.span
-        _apply_agent_finish_attributes(span, invocation)
-        otel_context.detach(invocation.context_token)
-        span.end()
-        return invocation
+        return self.stop(invocation)
 
-    def fail_agent(  # pylint: disable=no-self-use
+    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
-
-        span = invocation.span
-        _apply_agent_finish_attributes(span, invocation)
-        _apply_error_attributes(span, error)
-        otel_context.detach(invocation.context_token)
-        span.end()
-        return invocation
+        return self.fail(invocation, error)
 
     @contextmanager
     def agent(
@@ -291,42 +331,17 @@ def start_create_agent(
         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
+        return self.start(creation)
 
-    def stop_create_agent(self, creation: AgentCreation) -> AgentCreation:  # pylint: disable=no-self-use
+    def stop_create_agent(self, creation: AgentCreation) -> AgentCreation:
         """Finalize an agent creation successfully and end its span."""
-        if creation.context_token is None or creation.span is None:
-            return creation
-
-        span = creation.span
-        _apply_creation_finish_attributes(span, creation)
-        otel_context.detach(creation.context_token)
-        span.end()
-        return creation
+        return self.stop(creation)
 
-    def fail_create_agent(  # pylint: disable=no-self-use
+    def fail_create_agent(
         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
-
-        span = creation.span
-        _apply_creation_finish_attributes(span, creation)
-        _apply_error_attributes(span, error)
-        otel_context.detach(creation.context_token)
-        span.end()
-        return creation
+        return self.fail(creation, error)
 
     @contextmanager
     def create_agent(