common code for context manager

Author: lmolkova

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

@@ -16,10 +16,11 @@
 
 import timeit
 from abc import ABC, abstractmethod
+from contextlib import contextmanager
 from contextvars import Token
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Iterator
 
-from typing_extensions import TypeAlias
+from typing_extensions import Self, TypeAlias
 
 from opentelemetry._logs import Logger
 from opentelemetry.context import Context, attach, detach
@@ -127,3 +128,13 @@ def fail(self, error: Error | BaseException) -> None:
         if isinstance(error, BaseException):
             error = Error(type=type(error), message=str(error))
         self._finish(error)
+
+    @contextmanager
+    def _managed(self) -> Iterator[Self]:
+        """Context manager that calls stop() on success or fail() on exception."""
+        try:
+            yield self
+        except Exception as exc:
+            self.fail(exc)
+            raise
+        self.stop()

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

@@ -48,8 +48,7 @@
 
 from __future__ import annotations
 
-from contextlib import contextmanager
-from typing import Iterator
+from contextlib import AbstractContextManager
 
 from opentelemetry._logs import (
     LoggerProvider,
@@ -229,15 +228,14 @@ def fail_llm(  # pylint: disable=no-self-use
             invocation._inference_invocation.fail(error)
         return invocation
 
-    @contextmanager
     def inference(
         self,
         provider: str,
         *,
         request_model: str | None = None,
         server_address: str | None = None,
         server_port: int | None = None,
-    ) -> Iterator[InferenceInvocation]:
+    ) -> AbstractContextManager[InferenceInvocation]:
         """Context manager for LLM inference invocations.
 
         Only set data attributes on the invocation object, do not modify the span or context.
@@ -246,28 +244,21 @@ def inference(
         If an exception occurs inside the context, marks the span as error, ends it, and
         re-raises the original exception.
         """
-        invocation = self.start_inference(
+        return self.start_inference(
             provider=provider,
             request_model=request_model,
             server_address=server_address,
             server_port=server_port,
-        )
-        try:
-            yield invocation
-        except Exception as exc:
-            invocation.fail(exc)
-            raise
-        invocation.stop()
+        )._managed()
 
-    @contextmanager
     def embedding(
         self,
         provider: str,
         *,
         request_model: str | None = None,
         server_address: str | None = None,
         server_port: int | None = None,
-    ) -> Iterator[EmbeddingInvocation]:
+    ) -> AbstractContextManager[EmbeddingInvocation]:
         """Context manager for Embedding invocations.
 
         Only set data attributes on the invocation object, do not modify the span or context.
@@ -276,20 +267,13 @@ def embedding(
         If an exception occurs inside the context, marks the span as error, ends it, and
         re-raises the original exception.
         """
-        invocation = self.start_embedding(
+        return self.start_embedding(
             provider=provider,
             request_model=request_model,
             server_address=server_address,
             server_port=server_port,
-        )
-        try:
-            yield invocation
-        except Exception as exc:
-            invocation.fail(exc)
-            raise
-        invocation.stop()
+        )._managed()
 
-    @contextmanager
     def tool(
         self,
         name: str,
@@ -298,7 +282,7 @@ def tool(
         tool_call_id: str | None = None,
         tool_type: str | None = None,
         tool_description: str | None = None,
-    ) -> Iterator[ToolInvocation]:
+    ) -> AbstractContextManager[ToolInvocation]:
         """Context manager for Tool invocations.
 
         Only set data attributes on the invocation object, do not modify the span or context.
@@ -307,25 +291,18 @@ def tool(
         If an exception occurs inside the context, marks the span as error, ends it, and
         re-raises the original exception.
         """
-        invocation = self.start_tool(
+        return self.start_tool(
             name,
             arguments=arguments,
             tool_call_id=tool_call_id,
             tool_type=tool_type,
             tool_description=tool_description,
-        )
-        try:
-            yield invocation
-        except Exception as exc:
-            invocation.fail(exc)
-            raise
-        invocation.stop()
+        )._managed()
 
-    @contextmanager
     def workflow(
         self,
         name: str | None = None,
-    ) -> Iterator[WorkflowInvocation]:
+    ) -> AbstractContextManager[WorkflowInvocation]:
         """Context manager for Workflow invocations.
 
         Only set data attributes on the invocation object, do not modify the span or context.
@@ -334,14 +311,7 @@ def workflow(
         If an exception occurs inside the context, marks the span as error, ends it, and
         re-raises the original exception.
         """
-        invocation = self.start_workflow(name=name)
-
-        try:
-            yield invocation
-        except Exception as exc:
-            invocation.fail(exc)
-            raise
-        invocation.stop()
+        return self.start_workflow(name=name)._managed()
 
 
 def get_telemetry_handler(