clean up context handler, clarify unit tests

Author: keith-decker

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

@@ -98,7 +98,9 @@ def fail_llm(
         return invocation
 
     @contextmanager
-    def llm(self, invocation: LLMInvocation) -> Iterator[LLMInvocation]:
+    def llm(
+        self, invocation: Optional[LLMInvocation] = None
+    ) -> Iterator[LLMInvocation]:
         """Context manager for LLM invocations.
 
         Only set data attributes on the invocation object, do not modify the span or context.
@@ -107,6 +109,10 @@ def llm(self, invocation: LLMInvocation) -> Iterator[LLMInvocation]:
         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 = LLMInvocation(
+                request_model="",
+            )
         self.start_llm(invocation)
         try:
             yield invocation

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

@@ -48,7 +48,9 @@ def _apply_common_span_attributes(
     """
     request_model = invocation.request_model
     provider = invocation.provider
-
+    span.update_name(
+        f"{GenAI.GenAiOperationNameValues.CHAT.value} {request_model}"
+    )
     span.set_attribute(
         GenAI.GEN_AI_OPERATION_NAME, GenAI.GenAiOperationNameValues.CHAT.value
     )

util/opentelemetry-util-genai/tests/test_utils.py

@@ -135,14 +135,11 @@ def test_llm_start_and_stop_creates_span(self):  # pylint: disable=no-self-use
         )
 
         # Start and stop LLM invocation using context manager
-        invocation = LLMInvocation(
-            request_model="test-model",
-            input_messages=[message],
-            provider="test-provider",
-            attributes={"custom_attr": "value"},
-        )
-
-        with self.telemetry_handler.llm(invocation):
+        with self.telemetry_handler.llm() as invocation:
+            invocation.request_model = "test-model"
+            invocation.input_messages = [message]
+            invocation.provider = "test-provider"
+            invocation.attributes = {"custom_attr": "value"}
             assert invocation.span is not None
             invocation.output_messages = [chat_generation]
             invocation.attributes.update({"extra": "info"})
@@ -234,20 +231,16 @@ def test_parent_child_span_relationship(self):
             role="AI", parts=[Text(content="ok")], finish_reason="stop"
         )
 
-        # Start parent and child using nested contexts (child becomes child span of parent)
-        parent_invocation = LLMInvocation(
-            request_model="parent-model",
-            input_messages=[message],
-            provider="test-provider",
-        )
-        child_invocation = LLMInvocation(
-            request_model="child-model",
-            input_messages=[message],
-            provider="test-provider",
-        )
-
-        with self.telemetry_handler.llm(parent_invocation):
-            with self.telemetry_handler.llm(child_invocation):
+        with self.telemetry_handler.llm() as parent_invocation:
+            parent_invocation.request_model = "parent-model"
+            parent_invocation.input_messages = [message]
+            parent_invocation.provider = "test-provider"
+            # Perform things here, calling a tool, processing, etc.
+            with self.telemetry_handler.llm() as child_invocation:
+                child_invocation.request_model = "child-model"
+                child_invocation.input_messages = [message]
+                child_invocation.provider = "test-provider"
+                # Perform things here, calling a tool, processing, etc.
                 # Stop child first by exiting inner context
                 child_invocation.output_messages = [chat_generation]
             # Then stop parent by exiting outer context