add unit test for parent/child relationship

Author: keith-decker

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

@@ -230,6 +230,14 @@ def start(self, invocation: LLMInvocation):
         )
         if parent_state is not None:
             parent_state.children.append(invocation.run_id)
+        span = self._start_span(
+            name=f"{GenAI.GenAiOperationNameValues.CHAT.value} {invocation.request_model}",
+            kind=SpanKind.CLIENT,
+            parent_run_id=invocation.parent_run_id,
+        )
+        # Keep span active but do not end it here; it will be ended in finish()/error()
+        with use_span(span, end_on_exit=False):
+            self.spans[invocation.run_id] = _SpanState(span=span)
 
     @contextmanager
     def _start_span_for_invocation(self, invocation: LLMInvocation):
@@ -255,18 +263,39 @@ def _finalize_invocation(self, invocation: LLMInvocation) -> None:
         self._end_span(invocation.run_id)
 
     def finish(self, invocation: LLMInvocation):
-        with self._start_span_for_invocation(invocation) as span:
-            _apply_common_span_attributes(span, invocation)
-            _maybe_set_span_messages(
-                span, invocation.messages, invocation.chat_generations
-            )
+        state = self.spans.get(invocation.run_id)
+        if state is None:
+            with self._start_span_for_invocation(invocation) as span:
+                _apply_common_span_attributes(span, invocation)
+                _maybe_set_span_messages(
+                    span, invocation.messages, invocation.chat_generations
+                )
+            self._finalize_invocation(invocation)
+            return
+
+        span = state.span
+        _apply_common_span_attributes(span, invocation)
+        _maybe_set_span_messages(
+            span, invocation.messages, invocation.chat_generations
+        )
         self._finalize_invocation(invocation)
 
     def error(self, error: Error, invocation: LLMInvocation):
-        with self._start_span_for_invocation(invocation) as span:
-            span.set_status(Status(StatusCode.ERROR, error.message))
-            if span.is_recording():
-                span.set_attribute(
-                    ErrorAttributes.ERROR_TYPE, error.type.__qualname__
-                )
+        state = self.spans.get(invocation.run_id)
+        if state is None:
+            with self._start_span_for_invocation(invocation) as span:
+                span.set_status(Status(StatusCode.ERROR, error.message))
+                if span.is_recording():
+                    span.set_attribute(
+                        ErrorAttributes.ERROR_TYPE, error.type.__qualname__
+                    )
+            self._finalize_invocation(invocation)
+            return
+
+        span = state.span
+        span.set_status(Status(StatusCode.ERROR, error.message))
+        if span.is_recording():
+            span.set_attribute(
+                ErrorAttributes.ERROR_TYPE, error.type.__qualname__
+            )
         self._finalize_invocation(invocation)

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

@@ -169,3 +169,53 @@ def test_llm_start_and_stop_creates_span(self):  # pylint: disable=no-self-use
 
         assert isinstance(input_messages_json, str)
         assert isinstance(output_messages_json, str)
+
+    @patch_env_vars(
+        stability_mode="gen_ai_latest_experimental",
+        content_capturing="SPAN_ONLY",
+    )
+    def test_parent_child_span_relationship(self):
+        parent_id = uuid4()
+        child_id = uuid4()
+        message = InputMessage(role="Human", parts=[Text(content="hi")])
+        chat_generation = OutputMessage(
+            role="AI", parts=[Text(content="ok")], finish_reason="stop"
+        )
+
+        # Start parent and child (child references parent_run_id)
+        self.telemetry_handler.start_llm(
+            request_model="parent-model",
+            prompts=[message],
+            run_id=parent_id,
+            provider="test-provider",
+        )
+        self.telemetry_handler.start_llm(
+            request_model="child-model",
+            prompts=[message],
+            run_id=child_id,
+            parent_run_id=parent_id,
+            provider="test-provider",
+        )
+
+        # Stop child first, then parent (order should not matter)
+        self.telemetry_handler.stop_llm(
+            child_id, chat_generations=[chat_generation]
+        )
+        self.telemetry_handler.stop_llm(
+            parent_id, chat_generations=[chat_generation]
+        )
+
+        spans = self.span_exporter.get_finished_spans()
+        assert len(spans) == 2
+
+        # Identify spans irrespective of export order
+        child_span = next(s for s in spans if s.name == "chat child-model")
+        parent_span = next(s for s in spans if s.name == "chat parent-model")
+
+        # Same trace
+        assert child_span.context.trace_id == parent_span.context.trace_id
+        # Child has parent set to parent's span id
+        assert child_span.parent is not None
+        assert child_span.parent.span_id == parent_span.context.span_id
+        # Parent should not have a parent (root)
+        assert parent_span.parent is None