Merge branch 'main' into toolcall-handler-and-spans

Author: keith-decker

util/opentelemetry-util-genai/CHANGELOG.md

@@ -9,6 +9,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 - Add ToolCall span lifecycle support
   ([https://github.com/open-telemetry/opentelemetry-python-contrib/pull/4356/](#4356))
+- Add support for workflow in genAI utils handler.
+  ([https://github.com/open-telemetry/opentelemetry-python-contrib/pull/4366](#4366))
 - Enrich ToolCall type, breaking change: usage of ToolCall class renamed to ToolCallRequest 
   ([#4218](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/4218))
 - Add EmbeddingInvocation span lifecycle support

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

@@ -63,6 +63,7 @@
 
 from __future__ import annotations
 
+import logging
 import timeit
 from contextlib import contextmanager
 from typing import Iterator, TypeVar
@@ -88,10 +89,12 @@
     _apply_error_attributes,
     _apply_llm_finish_attributes,
     _apply_tool_call_attributes,
+    _apply_workflow_finish_attributes,
     _finish_tool_call_span,
     _get_embedding_span_name,
     _get_llm_span_name,
     _get_tool_call_span_name,
+    _get_workflow_span_name,
     _maybe_emit_llm_event,
 )
 from opentelemetry.util.genai.types import (
@@ -100,9 +103,27 @@
     GenAIInvocation,
     LLMInvocation,
     ToolCall,
+    WorkflowInvocation,
 )
 from opentelemetry.util.genai.version import __version__
 
+_logger = logging.getLogger(__name__)
+
+
+def _safe_detach(invocation: GenAIInvocation) -> None:
+    """Detach the context token if still present, as a safety net."""
+    if invocation.context_token is not None:
+        try:
+            otel_context.detach(invocation.context_token)
+        except Exception:  # pylint: disable=broad-except
+            pass
+    if invocation.span is not None:
+        try:
+            invocation.span.end()
+        except Exception:  # pylint: disable=broad-except
+            pass
+
+
 _T = TypeVar("_T", bound=GenAIInvocation)
 
 
@@ -158,20 +179,24 @@ def _record_metrics(
 
     def _start(self, invocation: _T) -> _T:
         """Start a GenAI invocation and create a pending span entry."""
-        span_kind = SpanKind.CLIENT
         if isinstance(invocation, LLMInvocation):
             span_name = _get_llm_span_name(invocation)
+            kind = SpanKind.CLIENT
         elif isinstance(invocation, EmbeddingInvocation):
             span_name = _get_embedding_span_name(invocation)
+            kind = SpanKind.CLIENT
         elif isinstance(invocation, ToolCall):
             span_name = _get_tool_call_span_name(invocation)
-            span_kind = SpanKind.INTERNAL
+            kind = SpanKind.INTERNAL
+        elif isinstance(invocation, WorkflowInvocation):
+            span_name = _get_workflow_span_name(invocation)
+            kind = SpanKind.INTERNAL
         else:
             span_name = ""
-
+            kind = SpanKind.CLIENT
         span = self._tracer.start_span(
             name=span_name,
-            kind=span_kind,
+            kind=kind,
         )
         if isinstance(invocation, ToolCall):
             _apply_tool_call_attributes(
@@ -203,6 +228,9 @@ def _stop(self, invocation: _T) -> _T:
             elif isinstance(invocation, ToolCall):
                 _finish_tool_call_span(span, invocation, capture_content=True)
                 self._record_metrics(invocation, span)
+            elif isinstance(invocation, WorkflowInvocation):
+                _apply_workflow_finish_attributes(span, invocation)
+                # TODO: Add workflow metrics when supported
         finally:
             # Detach context and end span even if finishing fails
             otel_context.detach(invocation.context_token)
@@ -234,6 +262,10 @@ def _fail(self, invocation: _T, error: Error) -> _T:
                 _finish_tool_call_span(span, invocation, capture_content=True)
                 self._record_metrics(invocation, span, error_type=error_type)
                 span.set_status(Status(StatusCode.ERROR, error.message))
+            elif isinstance(invocation, WorkflowInvocation):
+                _apply_workflow_finish_attributes(span, invocation)
+                _apply_error_attributes(span, error, error_type)
+                # TODO: Add workflow metrics when supported
         finally:
             # Detach context and end span even if finishing fails
             otel_context.detach(invocation.context_token)
@@ -347,6 +379,46 @@ def embedding(
             raise
         self.stop(invocation)
 
+    @contextmanager
+    def workflow(
+        self, invocation: WorkflowInvocation | None = None
+    ) -> Iterator[WorkflowInvocation]:
+        """Context manager for Workflow 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 = WorkflowInvocation()
+
+        try:
+            self.start(invocation)
+        except Exception:  # pylint: disable=broad-except
+            _logger.warning(
+                "Failed to start workflow telemetry", exc_info=True
+            )
+
+        try:
+            yield invocation
+        except Exception as exc:
+            try:
+                self.fail(invocation, Error(message=str(exc), type=type(exc)))
+            except Exception:  # pylint: disable=broad-except
+                _logger.warning(
+                    "Failed to record workflow failure", exc_info=True
+                )
+                _safe_detach(invocation)
+            raise
+
+        try:
+            self.stop(invocation)
+        except Exception:  # pylint: disable=broad-except
+            _logger.warning("Failed to stop workflow telemetry", exc_info=True)
+            _safe_detach(invocation)
+
 
 def get_telemetry_handler(
     tracer_provider: TracerProvider | None = None,

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

@@ -40,6 +40,7 @@
     MessagePart,
     OutputMessage,
     ToolCall,
+    WorkflowInvocation,
 )
 from opentelemetry.util.genai.utils import (
     ContentCapturingMode,
@@ -109,7 +110,14 @@ def _get_embedding_span_name(invocation: EmbeddingInvocation) -> str:
     return _get_span_name(invocation)
 
 
-def _get_llm_messages_attributes_for_span(
+def _get_workflow_span_name(invocation: WorkflowInvocation) -> str:
+    """Get the span name for an Workflow invocation."""
+    operation_name = invocation.operation_name
+    name = invocation.name
+    return f"{operation_name} {name}" if name else operation_name
+
+
+def _get_messages_attributes_for_span(
     input_messages: list[InputMessage],
     output_messages: list[OutputMessage],
     system_instruction: list[MessagePart] | None = None,
@@ -241,7 +249,7 @@ def _apply_llm_finish_attributes(
     attributes.update(_get_llm_request_attributes(invocation))
     attributes.update(_get_llm_response_attributes(invocation))
     attributes.update(
-        _get_llm_messages_attributes_for_span(
+        _get_messages_attributes_for_span(
             invocation.input_messages,
             invocation.output_messages,
             invocation.system_instruction,
@@ -346,6 +354,39 @@ def _get_llm_response_attributes(
     return {key: value for key, value in optional_attrs if value is not None}
 
 
+def _apply_workflow_finish_attributes(
+    span: Span, invocation: WorkflowInvocation
+) -> None:
+    """Apply attributes/messages common to finish() paths."""
+
+    # Build all attributes by reusing the attribute getter functions
+    attributes: dict[str, Any] = {}
+    attributes.update(_get_workflow_common_attributes(invocation))
+    attributes.update(
+        _get_messages_attributes_for_span(
+            invocation.input_messages,
+            invocation.output_messages,
+        )
+    )
+    attributes.update(invocation.attributes)
+
+    # Set all attributes on the span
+    if attributes:
+        span.set_attributes(attributes)
+
+
+def _get_workflow_common_attributes(
+    invocation: WorkflowInvocation,
+) -> dict[str, Any]:
+    """Get common Workflow attributes shared by finish() and error() paths.
+
+    Returns a dictionary of attributes.
+    """
+    return {
+        GenAI.GEN_AI_OPERATION_NAME: invocation.operation_name,
+    }
+
+
 def _get_embedding_response_attributes(
     invocation: EmbeddingInvocation,
 ) -> dict[str, Any]:
@@ -463,6 +504,8 @@ def _finish_tool_call_span(
     "_get_llm_response_attributes",
     "_get_llm_span_name",
     "_maybe_emit_llm_event",
+    "_get_workflow_common_attributes",
+    "_apply_workflow_finish_attributes",
     "_apply_embedding_finish_attributes",
     "_get_embedding_common_attributes",
     "_get_embedding_request_attributes",
@@ -471,4 +514,5 @@ def _finish_tool_call_span(
     "_apply_tool_call_attributes",
     "_finish_tool_call_span",
     "_get_tool_call_span_name",
+    "_get_workflow_span_name",
 ]

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

@@ -285,6 +285,9 @@ class WorkflowInvocation(GenAIInvocation):
         default_factory=_new_output_messages
     )
 
+    def __post_init__(self) -> None:
+        self.operation_name = "invoke_workflow"
+
 
 @dataclass
 class LLMInvocation(GenAIInvocation):