GenAI Utils | Allow passing external Span to the manual LLMInvocation

Author: Fiery-Fenix

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

@@ -137,16 +137,20 @@ def start_llm(
     ) -> LLMInvocation:
         """Start an LLM invocation and create a pending span entry."""
         # Create a span and attach it as current; keep the token to detach later
-        span = self._tracer.start_span(
-            name=f"{invocation.operation_name} {invocation.request_model}",
-            kind=SpanKind.CLIENT,
-        )
+        # If the span is already set on the invocation - use it instead of creating a new one
+        # This allows users to control the span outside of the handler if needed
+        if invocation.span is None:
+            span = self._tracer.start_span(
+                name=f"{invocation.operation_name} {invocation.request_model}",
+                kind=SpanKind.CLIENT,
+            )
+            invocation.span = span
+            invocation.end_span_on_exit = True
         # Record a monotonic start timestamp (seconds) for duration
         # calculation using timeit.default_timer.
         invocation.monotonic_start_s = timeit.default_timer()
-        invocation.span = span
         invocation.context_token = otel_context.attach(
-            set_span_in_context(span)
+            set_span_in_context(invocation.span)
         )
         return invocation
 
@@ -162,7 +166,9 @@ def stop_llm(self, invocation: LLMInvocation) -> LLMInvocation:  # pylint: disab
         _maybe_emit_llm_event(self._logger, span, invocation)
         # Detach context and end span
         otel_context.detach(invocation.context_token)
-        span.end()
+        # End the span only if it was created by the handler
+        if invocation.end_span_on_exit:
+            span.end()
         return invocation
 
     def fail_llm(  # pylint: disable=no-self-use
@@ -181,7 +187,9 @@ def fail_llm(  # pylint: disable=no-self-use
         _maybe_emit_llm_event(self._logger, span, invocation, error)
         # Detach context and end span
         otel_context.detach(invocation.context_token)
-        span.end()
+        # End the span only if it was created by the handler
+        if invocation.end_span_on_exit:
+            span.end()
         return invocation
 
     @contextmanager

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

@@ -201,6 +201,7 @@ class GenAIInvocation:
     context_token: ContextToken | None = None
     span: Span | None = None
     attributes: dict[str, Any] = field(default_factory=_new_str_any_dict)
+    end_span_on_exit: bool = False
 
 
 @dataclass

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

@@ -861,6 +861,59 @@ def test_emits_llm_event_by_default_for_span_and_event(self):
         )
         self.assertIn(GenAI.GEN_AI_INPUT_MESSAGES, log_record.attributes)
 
+    @patch_env_vars(
+        stability_mode="gen_ai_latest_experimental",
+        content_capturing="SPAN_ONLY",
+        emit_event="",
+    )
+    def test_llm_manual_start_and_stop_external_span(self):
+        message = _create_input_message("hi")
+        chat_generation = _create_output_message("ok")
+
+        # Create a custom Span
+        external_span = self.telemetry_handler._tracer.start_span(
+            name="external operation", kind=trace.SpanKind.INTERNAL
+        )
+
+        invocation = LLMInvocation(
+            span=external_span,
+            request_model="manual-model",
+            input_messages=[message],
+            provider="test-provider",
+            attributes={"external": True},
+        )
+
+        self.telemetry_handler.start_llm(invocation)
+        assert invocation.span is not None
+        assert invocation.span == external_span
+        assert invocation.end_span_on_exit is False
+
+        invocation.output_messages = [chat_generation]
+        self.telemetry_handler.stop_llm(invocation)
+        # Verify that the external span is still recording after stop_llm is called
+        assert invocation.span.is_recording()
+        # Manually end the external span after stopping the LLM invocation
+        external_span.end()
+
+        span = _get_single_span(self.span_exporter)
+        # Make sure that external span is used and not overwritten by telemetry handler
+        assert span.kind == trace.SpanKind.INTERNAL
+        _assert_span_time_order(span)
+
+        attrs = _get_span_attributes(span)
+        _assert_span_attributes(
+            attrs,
+            {
+                GenAI.GEN_AI_OPERATION_NAME: "chat",
+                GenAI.GEN_AI_REQUEST_MODEL: "manual-model",
+                GenAI.GEN_AI_PROVIDER_NAME: "test-provider",
+                GenAI.GEN_AI_INPUT_MESSAGES: AnyNonNone(),
+                GenAI.GEN_AI_OUTPUT_MESSAGES: AnyNonNone(),
+                GenAI.GEN_AI_RESPONSE_FINISH_REASONS: ("stop",),
+                "external": True,
+            },
+        )
+
 
 class AnyNonNone:
     def __eq__(self, other):