GenAI Utils minor API cleanup

Author: lmolkova

tox.ini

@@ -790,6 +790,7 @@ deps =
   util-genai: {[testenv]test_deps}
   util-genai: -r {toxinidir}/util/opentelemetry-util-genai/test-requirements.txt
   util-genai: {toxinidir}/util/opentelemetry-util-genai
+  lint-util-genai: {toxinidir}/util/opentelemetry-util-genai
 
   ; FIXME: add coverage testing
 allowlist_externals =

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

@@ -28,9 +28,13 @@
 Usage:
     handler = get_telemetry_handler()
 
-    # Create an invocation object with your request data
-    # The span and context_token attributes are set by the TelemetryHandler, and
-    # managed by the TelemetryHandler during the lifecycle of the span.
+    # Create an invocation object
+    invocation = LLMInvocation(
+        request_model="my-model",
+        input_messages=[...],
+        provider="my-provider"
+    )
+
 
     # Use the context manager to manage the lifecycle of an LLM invocation.
     with handler.llm(invocation) as invocation:
@@ -39,31 +43,24 @@
         invocation.attributes.update({"more": "attrs"})
 
     # Or, if you prefer to manage the lifecycle manually
-    invocation = LLMInvocation(
-        request_model="my-model",
-        input_messages=[...],
-        provider="my-provider",
-        attributes={"custom": "attr"},
-    )
-
-    # Start the invocation (opens a span)
     handler.start_llm(invocation)
 
     # Populate outputs and any additional attributes, then stop (closes the span)
     invocation.output_messages = [...]
     invocation.attributes.update({"more": "attrs"})
     handler.stop_llm(invocation)
 
-    # Or, in case of error
-    handler.fail_llm(invocation, Error(type="...", message="..."))
+    # Or, in case of error — pass an exception directly or an Error object
+    handler.fail_llm(invocation, exc)
+    handler.fail_llm(invocation, Error(type=MyError, message="..."))
 """
 
 from __future__ import annotations
 
 import logging
 import timeit
 from contextlib import contextmanager
-from typing import Iterator, TypeVar
+from typing import Iterator, TypeVar, overload
 
 from opentelemetry import context as otel_context
 from opentelemetry._logs import (
@@ -104,9 +101,9 @@
 
 def _safe_detach(invocation: GenAIInvocation) -> None:
     """Detach the context token if still present, as a safety net."""
-    if invocation.context_token is not None:
+    if invocation._context_token is not None:
         try:
-            otel_context.detach(invocation.context_token)
+            otel_context.detach(invocation._context_token)
         except Exception:  # pylint: disable=broad-except
             pass
     if invocation.span is not None:
@@ -197,16 +194,16 @@ def _start(self, invocation: _T) -> _T:
         )
         # Record a monotonic start timestamp (seconds) for duration
         # calculation using timeit.default_timer.
-        invocation.monotonic_start_s = timeit.default_timer()
+        invocation._monotonic_start_s = timeit.default_timer()
         invocation.span = span
-        invocation.context_token = otel_context.attach(
+        invocation._context_token = otel_context.attach(
             set_span_in_context(span)
         )
         return invocation
 
     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:
+        if invocation._context_token is None or invocation.span is None:
             # TODO: Provide feedback that this invocation was not started
             return invocation
 
@@ -224,13 +221,13 @@ def _stop(self, invocation: _T) -> _T:
                 # TODO: Add workflow metrics when supported
         finally:
             # Detach context and end span even if finishing fails
-            otel_context.detach(invocation.context_token)
+            otel_context.detach(invocation._context_token)
             span.end()
         return invocation
 
     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:
+        if invocation._context_token is None or invocation.span is None:
             # TODO: Provide feedback that this invocation was not started
             return invocation
 
@@ -258,7 +255,7 @@ def _fail(self, invocation: _T, error: Error) -> _T:
                 # TODO: Add workflow metrics when supported
         finally:
             # Detach context and end span even if finishing fails
-            otel_context.detach(invocation.context_token)
+            otel_context.detach(invocation._context_token)
             span.end()
         return invocation
 
@@ -273,8 +270,16 @@ 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:
+    @overload
+    def fail(self, invocation: _T, error: Error) -> _T: ...
+
+    @overload
+    def fail(self, invocation: _T, error: BaseException) -> _T: ...
+
+    def fail(self, invocation: _T, error: Error | BaseException) -> _T:
         """Fail a GenAI invocation and end its span with error status."""
+        if isinstance(error, BaseException):
+            error = Error(type=type(error), message=str(error))
         return self._fail(invocation, error)
 
     # LLM-specific convenience methods
@@ -286,10 +291,22 @@ def stop_llm(self, invocation: LLMInvocation) -> LLMInvocation:
         """Finalize an LLM invocation successfully and end its span."""
         return self._stop(invocation)
 
+    @overload
     def fail_llm(
         self, invocation: LLMInvocation, error: Error
+    ) -> LLMInvocation: ...
+
+    @overload
+    def fail_llm(
+        self, invocation: LLMInvocation, error: BaseException
+    ) -> LLMInvocation: ...
+
+    def fail_llm(
+        self, invocation: LLMInvocation, error: Error | BaseException
     ) -> LLMInvocation:
         """Fail an LLM invocation and end its span with error status."""
+        if isinstance(error, BaseException):
+            error = Error(type=type(error), message=str(error))
         return self._fail(invocation, error)
 
     @contextmanager
@@ -312,7 +329,7 @@ def llm(
         try:
             yield invocation
         except Exception as exc:
-            self.fail_llm(invocation, Error(message=str(exc), type=type(exc)))
+            self.fail_llm(invocation, exc)
             raise
         self.stop_llm(invocation)
 
@@ -334,7 +351,7 @@ def embedding(
         try:
             yield invocation
         except Exception as exc:
-            self.fail(invocation, Error(message=str(exc), type=type(exc)))
+            self.fail(invocation, exc)
             raise
         self.stop(invocation)
 
@@ -364,7 +381,7 @@ def workflow(
             yield invocation
         except Exception as exc:
             try:
-                self.fail(invocation, Error(message=str(exc), type=type(exc)))
+                self.fail(invocation, exc)
             except Exception:  # pylint: disable=broad-except
                 _logger.warning(
                     "Failed to record workflow failure", exc_info=True

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

@@ -81,9 +81,9 @@ def record(
 
         # Calculate duration from span timing or invocation monotonic start
         duration_seconds: Optional[float] = None
-        if invocation.monotonic_start_s is not None:
+        if invocation._monotonic_start_s is not None:
             duration_seconds = max(
-                timeit.default_timer() - invocation.monotonic_start_s,
+                timeit.default_timer() - invocation._monotonic_start_s,
                 0.0,
             )
 

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

@@ -67,7 +67,7 @@ def _get_llm_common_attributes(
 
     return {
         GenAI.GEN_AI_OPERATION_NAME: invocation.operation_name,
-        **{key: value for key, value in optional_attrs if value is not None},
+        **{key: value for key, value in optional_attrs if value},
     }
 
 
@@ -79,14 +79,14 @@ def _get_embedding_common_attributes(
     Returns a dictionary of attributes.
     """
     optional_attrs = (
+        (GenAI.GEN_AI_PROVIDER_NAME, invocation.provider),
         (server_attributes.SERVER_ADDRESS, invocation.server_address),
         (server_attributes.SERVER_PORT, invocation.server_port),
     )
 
     return {
         GenAI.GEN_AI_OPERATION_NAME: invocation.operation_name,
-        GenAI.GEN_AI_PROVIDER_NAME: invocation.provider,
-        **{key: value for key, value in optional_attrs if value is not None},
+        **{key: value for key, value in optional_attrs if value},
     }
 
 

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

@@ -256,24 +256,42 @@ def _new_str_any_dict() -> dict[str, Any]:
 
 @dataclass
 class GenAIInvocation:
-    context_token: ContextToken | None = None
-    span: Span | None = None
+    """
+    Base class for all GenAI invocation types. Manages the lifecycle of a single
+    GenAI operation (LLM call, embedding, tool execution, workflow, etc.).
+    """
+
+    span: Span | None = field(default=None, init=False)
+    _context_token: ContextToken | None = field(default=None, init=False)
+    _monotonic_start_s: float | None = field(default=None, init=False)
+
+    operation_name: str = ""
+    """
+    GenAI operation name (maps to gen_ai.operation.name). Subclasses hardcode
+    this to one of the GenAiOperationNameValues enum values. Callers may
+    override it, but this is not recommended.
+    """
+
     attributes: dict[str, Any] = field(default_factory=_new_str_any_dict)
-    error_type: str | None = None
+    """
+    Additional attributes to set on spans and/or events. Not set on metrics.
+    """
 
-    monotonic_start_s: float | None = None
+    metric_attributes: dict[str, Any] = field(
+        default_factory=_new_str_any_dict
+    )
     """
-    Monotonic start time in seconds (from timeit.default_timer) used for
-    duration calculations to avoid mixing clock sources. This is populated
-    by the TelemetryHandler when starting an invocation.
+    Additional attributes to set on metrics. Must be low cardinality.
+    Not set on spans or events.
     """
 
 
 @dataclass
 class WorkflowInvocation(GenAIInvocation):
     """
-    Represents predetermined static sequence of operations eg: Agent, LLM, tool, and retrieval invocations.
-    A workflow groups multiple operations together, accepting input(s) and producing final output(s).
+    Represents a predetermined sequence of operations (e.g. agent, LLM, tool,
+    and retrieval invocations). A workflow groups multiple operations together,
+    accepting input(s) and producing final output(s).
     """
 
     name: str = ""
@@ -291,11 +309,7 @@ def __post_init__(self) -> None:
 
 @dataclass
 class LLMInvocation(GenAIInvocation):
-    """
-    Represents a single LLM call invocation. When creating an LLMInvocation object,
-    only update the data attributes. The span and context_token attributes are
-    set by the TelemetryHandler.
-    """
+    """Represents a single LLM chat/completion call."""
 
     operation_name: str = GenAI.GenAiOperationNameValues.CHAT.value
     request_model: str | None = None
@@ -308,24 +322,12 @@ class LLMInvocation(GenAIInvocation):
     system_instruction: list[MessagePart] = field(
         default_factory=_new_system_instruction
     )
-    provider: str | None = None
+    provider: str = ""
     response_model_name: str | None = None
     response_id: str | None = None
     finish_reasons: list[str] | None = None
     input_tokens: int | None = None
     output_tokens: int | None = None
-    attributes: dict[str, Any] = field(default_factory=_new_str_any_dict)
-    """
-    Additional attributes to set on spans and/or events. These attributes
-    will not be set on metrics.
-    """
-    metric_attributes: dict[str, Any] = field(
-        default_factory=_new_str_any_dict
-    )
-    """
-    Additional attributes to set on metrics. Must be of a low cardinality.
-    These attributes will not be set on spans or events.
-    """
     temperature: float | None = None
     top_p: float | None = None
     frequency_penalty: float | None = None
@@ -339,15 +341,11 @@ class LLMInvocation(GenAIInvocation):
 
 @dataclass
 class EmbeddingInvocation(GenAIInvocation):
-    """
-    Represents a single embedding model invocation. When creating an
-    EmbeddingInvocation object, only update the data attributes. The span
-    and context_token attributes are set by the TelemetryHandler.
-    """
+    """Represents a single embedding model invocation."""
 
     operation_name: str = GenAI.GenAiOperationNameValues.EMBEDDINGS.value
     request_model: str | None = None
-    provider: str | None = None  # e.g., azure.ai.openai, openai, aws.bedrock
+    provider: str = ""  # e.g., azure.ai.openai, openai, aws.bedrock
     server_address: str | None = None
     server_port: int | None = None
 
@@ -358,33 +356,12 @@ class EmbeddingInvocation(GenAIInvocation):
     dimension_count: int | None = None
     response_model_name: str | None = None
 
-    attributes: dict[str, Any] = field(default_factory=_new_str_any_dict)
-    """
-    Additional attributes to set on spans and/or events. These attributes
-    will not be set on metrics.
-    """
-
-    metric_attributes: dict[str, Any] = field(
-        default_factory=_new_str_any_dict
-    )
-    """
-    Additional attributes to set on metrics. Must be of a low cardinality.
-    These attributes will not be set on spans or events.
-    """
-
 
 @dataclass()
 class ToolCall(GenAIInvocation):
-    """Represents a tool call for execution tracking with spans and metrics.
+    """Represents a tool call invocation for execute_tool span tracking.
 
-    This type extends GenAIInvocation (like LLMInvocation) for consistent lifecycle
-    management across all invocation types. It is NOT used as a MessagePart directly -
-    use ToolCallRequest for that purpose.
-
-    Inherits from GenAIInvocation:
-    - context_token: Context tracking for span lifecycle
-    - span: Active span reference
-    - attributes: Custom attributes dict for extensibility
+    Not used as a message part — use ToolCallRequest for that purpose.
 
     Reference: https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-spans.md#execute-tool-span
 
@@ -414,9 +391,6 @@ class ToolCall(GenAIInvocation):
     # gen_ai.tool.call.result - Result returned by the tool (Opt-In, may contain sensitive data)
     tool_result: Any = None
 
-    # Timing field (not inherited from GenAIInvocation, matches LLMInvocation pattern)
-    monotonic_start_s: float | None = None
-
 
 @dataclass
 class Error: