Refactor GenAIInvocation to be a ContextManager to simplify code. Minor change to ToolInvocation. (#17)

* Make a few minor changes to utils.. Remove start_tool in favor of just using tool

* Respond to copilot comments

* Fix typecheck

* Respond to PR comments

* Fix changelog and lint issue

* Fix typecheck

* Fix typing..

---------

Co-authored-by: Liudmila Molkova <neskazu@gmail.com>

Author: DylanRussell

util/opentelemetry-util-genai/.changelog/17.added

@@ -0,0 +1 @@
+Add `tool_result` as a parameter to `ToolInvocation`, add `tool_result` and `arguments` as span attributes to the `execute_tool` span if `ContentCapture` flag is set.

util/opentelemetry-util-genai/.changelog/17.deprecated

@@ -0,0 +1 @@
+deprecate all `start_` factories, update all `invocation` factories to return objects that can be used as ContextManager's

util/opentelemetry-util-genai/AGENTS.md

@@ -25,7 +25,7 @@ applicable ones.
 Every new operation type must follow this pattern:
 
 ```python
-invocation = handler.start_inference(provider, request_model, server_address=..., server_port=...)
+invocation = handler.inference(provider, request_model, server_address=..., server_port=...)
 invocation.temperature = ...
 try:
     response = client.call(...)
@@ -39,24 +39,23 @@ except Exception as exc:
 
 Factory methods on `TelemetryHandler` (`handler.py`):
 
-- `start_inference(provider, request_model, *, server_address, server_port)` → `InferenceInvocation`
-- `start_embedding(provider, request_model, *, server_address, server_port)` → `EmbeddingInvocation`
-- `start_tool(name, *, arguments, tool_call_id, tool_type, tool_description)` → `ToolInvocation`
-- `start_workflow(name)` → `WorkflowInvocation`
+- `inference(provider, request_model, *, server_address, server_port)` → `InferenceInvocation`
+- `embedding(provider, request_model, *, server_address, server_port)` → `EmbeddingInvocation`
+- `tool(name, *, arguments, tool_call_id, tool_type, tool_description)` → `ToolInvocation`
+- `workflow(name)` → `WorkflowInvocation`
 
-Context manager equivalents (`handler.inference()`, `handler.embedding()`, `handler.tool()`,
-`handler.workflow()`) are available when the span lifetime maps cleanly to a `with` block.
+The returned object can also be used as a context manager (`with ... as invocation:`) when the span lifetime maps cleanly to a `with` block.
 
-`start_*()` factories must map 1:1 to distinct semconv operation types (inference, embeddings,
+The above factories must map 1:1 to distinct semconv operation types (inference, embeddings,
 tool execution, agent invocation, workflow invocation). Names must match the operation
 unambiguously — for example, `create_agent` and `invoke_agent` are different operations, so a
-single `start_agent()` would be ambiguous and is not acceptable. Add a new factory per operation
+single `agent()` would be ambiguous and is not acceptable. Add a new factory per operation
 instead.
 
-Factory names are Python-style singular verbs (`start_embedding`, `start_tool`); the op names
-they map to follow semconv (`embeddings`, `tool execution`, future operations).
+Factory names are Python-style singular verbs (`inference`, `embedding`, `tool`, `workflow`); the op names
+they map to follow semconv operations.
 
-`start_*()` factories must accept all attributes that semconv marks as important for sampling
+Factory methods must accept all attributes that semconv marks as important for sampling
 decisions as parameters, so they are on the span at creation time. Attributes that are also
 marked required by semconv must be required parameters (no default value). Operation name
 is usually hardcoded in specific invocation and does not need to be passed.
@@ -65,7 +64,7 @@ is usually hardcoded in specific invocation and does not need to be passed.
 
 **Never construct invocation types directly** (`InferenceInvocation(...)`, `ToolInvocation(...)`,
 etc.) in instrumentation or production code — direct construction skips span creation and context
-propagation, so all telemetry calls become no-ops. Always use `handler.start_*()`.
+propagation, so all telemetry calls become no-ops. Always use `handler.*()`.
 
 ## 3. Exception Handling
 

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

@@ -36,9 +36,8 @@
 class AgentInvocation(GenAIInvocation):
     """Represents a single agent invocation (invoke_agent span).
 
-    Use handler.start_invoke_local_agent() / handler.start_invoke_remote_agent()
-    or the handler.invoke_local_agent() / handler.invoke_remote_agent() context
-    managers rather than constructing this directly.
+    Use handler.invoke_local_agent() or handler.invoke_remote_agent()
+    rather than constructing this directly.
 
     Reference:
         Client span: https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/gen-ai-agent-spans.md#invoke-agent-client-span
@@ -59,7 +58,7 @@ def __init__(
         server_port: int | None = None,
         agent_name: str | None = None,
     ) -> None:
-        """Use handler.start_invoke_local_agent() or handler.start_invoke_remote_agent() instead of calling this directly."""
+        """Use handler.invoke_local_agent() or handler.invoke_remote_agent() instead of calling this directly."""
         _operation_name = GenAI.GenAiOperationNameValues.INVOKE_AGENT.value
         super().__init__(
             tracer,

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

@@ -20,8 +20,7 @@
 class EmbeddingInvocation(GenAIInvocation):
     """Represents a single embedding model invocation.
 
-    Use handler.start_embedding(provider) or the handler.embedding(provider)
-    context manager rather than constructing this directly.
+    Use handler.embedding(provider) rather than constructing this directly.
     """
 
     def __init__(
@@ -36,7 +35,7 @@ def __init__(
         server_address: str | None = None,
         server_port: int | None = None,
     ) -> None:
-        """Use handler.start_embedding(provider) or handler.embedding(provider) instead of calling this directly."""
+        """Use handler.embedding(provider) rather than calling this directly."""
         _operation_name = GenAI.GenAiOperationNameValues.EMBEDDINGS.value
         super().__init__(
             tracer,

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

@@ -37,8 +37,7 @@
 class InferenceInvocation(GenAIInvocation):
     """Represents a single LLM chat/completion call.
 
-    Use handler.start_inference(provider) or the handler.inference(provider)
-    context manager rather than constructing this directly.
+    Use handler.inference(provider) rather than constructing this directly.
     """
 
     def __init__(
@@ -57,7 +56,7 @@ def __init__(
         operation_name = (
             operation_name or GenAI.GenAiOperationNameValues.CHAT.value
         )
-        """Use handler.start_inference(provider) or handler.inference(provider) instead of calling this directly."""
+        """Use handler.inference(provider) rather than calling this directly."""
         super().__init__(
             tracer,
             metrics_recorder,

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

@@ -4,13 +4,14 @@
 from __future__ import annotations
 
 import timeit
-from abc import ABC, abstractmethod
-from contextlib import contextmanager
+from abc import abstractmethod
+from contextlib import AbstractContextManager
 from contextvars import Token
 from dataclasses import asdict
-from typing import TYPE_CHECKING, Any, Generator, Sequence
+from types import TracebackType
+from typing import TYPE_CHECKING, Any, Sequence
 
-from typing_extensions import Self, TypeAlias
+from typing_extensions import TypeAlias
 
 from opentelemetry._logs import Logger, LogRecord
 from opentelemetry.context import Context, attach, detach
@@ -39,10 +40,11 @@
 if TYPE_CHECKING:
     from opentelemetry.util.genai.metrics import InvocationMetricsRecorder
 
+
 ContextToken: TypeAlias = Token[Context]
 
 
-class GenAIInvocation(ABC):
+class GenAIInvocation(AbstractContextManager["GenAIInvocation"]):
     """
     Base class for all GenAI invocation types. Manages the lifecycle of a single
     GenAI operation (LLM call, embedding, tool execution, workflow, etc.).
@@ -165,15 +167,19 @@ def fail(self, error: Error | BaseException) -> None:
             error = Error(type=type(error), message=str(error))
         self._finish(error)
 
-    @contextmanager
-    def _managed(self) -> Generator[Self, None, None]:
-        """Context manager that calls stop() on success or fail() on exception."""
-        try:
-            yield self
-        except Exception as exc:
-            self.fail(exc)
-            raise
-        self.stop()
+    def __enter__(self) -> GenAIInvocation:
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        if exc_value is not None and isinstance(exc_value, Exception):
+            self.fail(exc_value)
+        else:
+            self.stop()
 
 
 def get_content_attributes(

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

@@ -13,6 +13,10 @@
 from opentelemetry.util.genai._invocation import Error, GenAIInvocation
 from opentelemetry.util.genai.completion_hook import CompletionHook
 from opentelemetry.util.genai.metrics import InvocationMetricsRecorder
+from opentelemetry.util.genai.utils import (
+    should_capture_content_on_spans,
+)
+from opentelemetry.util.types import AttributeValue
 
 
 class ToolInvocation(GenAIInvocation):
@@ -43,10 +47,11 @@ def __init__(
         completion_hook: CompletionHook,
         name: str,
         *,
-        arguments: Any = None,
+        arguments: AttributeValue | None = None,
         tool_call_id: str | None = None,
         tool_type: str | None = None,
         tool_description: str | None = None,
+        tool_result: AttributeValue | None = None,
     ) -> None:
         """Use handler.start_tool(name) or handler.tool(name) instead of calling this directly."""
         _operation_name = GenAI.GenAiOperationNameValues.EXECUTE_TOOL.value
@@ -58,12 +63,13 @@ def __init__(
             operation_name=_operation_name,
             span_name=f"{_operation_name} {name}" if name else _operation_name,
         )
+        self.should_capture_content_on_span = should_capture_content_on_spans()
         self.name = name
+        self.tool_result = tool_result
         self.arguments = arguments
         self.tool_call_id = tool_call_id
         self.tool_type = tool_type
         self.tool_description = tool_description
-        self.tool_result: Any = None
         self._start(self._get_base_attributes())
 
     def _get_base_attributes(self) -> dict[str, Any]:
@@ -73,6 +79,18 @@ def _get_base_attributes(self) -> dict[str, Any]:
             (GenAI.GEN_AI_TOOL_CALL_ID, self.tool_call_id),
             (GenAI.GEN_AI_TOOL_TYPE, self.tool_type),
             (GenAI.GEN_AI_TOOL_DESCRIPTION, self.tool_description),
+            (
+                GenAI.GEN_AI_TOOL_CALL_ARGUMENTS,
+                self.arguments
+                if self.should_capture_content_on_span
+                else None,
+            ),
+            (
+                GenAI.GEN_AI_TOOL_CALL_RESULT,
+                self.tool_result
+                if self.should_capture_content_on_span
+                else None,
+            ),
         )
         return {
             GenAI.GEN_AI_OPERATION_NAME: self._operation_name,
@@ -94,7 +112,18 @@ def _apply_finish(self, error: Error | None = None) -> None:
             (GenAI.GEN_AI_TOOL_CALL_ID, self.tool_call_id),
             (GenAI.GEN_AI_TOOL_TYPE, self.tool_type),
             (GenAI.GEN_AI_TOOL_DESCRIPTION, self.tool_description),
-            (GenAI.GEN_AI_TOOL_CALL_ARGUMENTS, self.arguments),
+            (
+                GenAI.GEN_AI_TOOL_CALL_ARGUMENTS,
+                self.arguments
+                if self.should_capture_content_on_span
+                else None,
+            ),
+            (
+                GenAI.GEN_AI_TOOL_CALL_RESULT,
+                self.tool_result
+                if self.should_capture_content_on_span
+                else None,
+            ),
         )
         attributes: dict[str, Any] = {
             GenAI.GEN_AI_OPERATION_NAME: self._operation_name,

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

@@ -32,8 +32,7 @@ class WorkflowInvocation(GenAIInvocation):
     and retrieval invocations). A workflow groups multiple operations together,
     accepting input(s) and producing final output(s).
 
-    Use handler.start_workflow(name) or the handler.workflow(name) context
-    manager rather than constructing this directly.
+    Use handler.workflow(name) rather than constructing this directly.
     """
 
     def __init__(
@@ -44,7 +43,7 @@ def __init__(
         completion_hook: CompletionHook,
         name: str | None,
     ) -> None:
-        """Use handler.start_workflow(name) or handler.workflow(name) instead of calling this directly."""
+        """Use handler.workflow(name) rather than calling this directly."""
         _operation_name = "invoke_workflow"
         super().__init__(
             tracer,