Merge branch 'main' into fix/azure-openai-responses-structured-predict-engine

Author: joshuakilts

llama-index-integrations/observability/llama-index-observability-otel/llama_index/observability/otel/base.py

@@ -1,5 +1,6 @@
 import inspect
 import logging
+from contextvars import Token
 from datetime import datetime
 from typing import Any, Dict, List, Literal, Mapping, Optional, Sequence, Union, cast
 
@@ -52,6 +53,9 @@ class OTelCompatibleSpanHandler(SimpleSpanHandler):
     _events_by_span: Dict[str, List[OTelEventAttributes]] = PrivateAttr(
         default_factory=dict,
     )
+    _context_tokens: Dict[str, Token[context.Context]] = PrivateAttr(
+        default_factory=dict
+    )
     all_spans: Dict[str, trace.Span] = Field(
         default_factory=dict, description="All the registered OpenTelemetry spans."
     )
@@ -79,6 +83,7 @@ def __init__(
         self._tracer = tracer
         self._tracer_provider = tracer_provider
         self._events_by_span = {}
+        self._context_tokens = {}
         self.debug = debug
 
     def close(self) -> None:
@@ -117,6 +122,32 @@ def restore_propagation_context(self, ctx: Dict[str, Any]) -> None:
             restored = propagate.extract(carrier)
             context.attach(restored)
 
+    def _get_parent_context(self, parent_span_id: Optional[str]) -> context.Context:
+        ctx = context.get_current()
+        current_span_context = trace.get_current_span(ctx).get_span_context()
+        if (
+            current_span_context.is_valid
+            or parent_span_id is None
+            or parent_span_id not in self.all_spans
+        ):
+            return ctx
+        return set_span_in_context(span=self.all_spans[parent_span_id])
+
+    def _add_events_to_span(self, id_: str, span: trace.Span) -> None:
+        events = self._events_by_span.pop(id_, [])
+        for event in events:
+            span.add_event(name=event.name, attributes=event.attributes)
+
+    def _detach_context_and_end_span(self, id_: str, span: trace.Span) -> None:
+        token = self._context_tokens.pop(id_, None)
+        try:
+            if token is not None:
+                context.detach(token)
+        except BaseException:
+            _logger.warning("Error detaching OTel context for %s", id_, exc_info=True)
+        finally:
+            span.end()
+
     def new_span(
         self,
         id_: str,
@@ -137,14 +168,14 @@ def new_span(
         # 1. Same-process parent found in all_spans → use it
         # 2. Otherwise → use ambient OTel context (set by restore_propagation_context
         #    for cross-process, or inherited naturally for same-process roots)
-        if parent_span_id is not None and parent_span_id in self.all_spans:
-            ctx = set_span_in_context(span=self.all_spans[parent_span_id])
-        else:
-            ctx = context.get_current()
+        ctx = self._get_parent_context(parent_span_id)
 
         otel_span = self._tracer.start_span(name=span_name, context=ctx)
         self.all_spans.update({id_: otel_span})
 
+        token = context.attach(set_span_in_context(otel_span))
+        self._context_tokens[id_] = token
+
         # Record instrument_tags as span attributes
         if tags is not None:
             for key, value in tags.items():
@@ -180,13 +211,9 @@ def prepare_to_exit_span(
             _logger.warning("No OTel span found for %s in prepare_to_exit_span", id_)
             return sp
 
-        # Get and process events specific to this span
-        events = self._events_by_span.pop(id_, [])
-        for event in events:
-            span.add_event(name=event.name, attributes=event.attributes)
-
+        self._add_events_to_span(id_, span)
         span.set_status(status=trace.StatusCode.OK)
-        span.end()
+        self._detach_context_and_end_span(id_, span)
         return sp
 
     def prepare_to_drop_span(
@@ -209,15 +236,11 @@ def prepare_to_drop_span(
             _logger.warning("No OTel span found for %s in prepare_to_drop_span", id_)
             return sp
 
-        # Get and process events specific to this span
-        events = self._events_by_span.pop(id_, [])
-        for event in events:
-            span.add_event(name=event.name, attributes=event.attributes)
-
+        self._add_events_to_span(id_, span)
         if err is not None:
             span.record_exception(err)
         span.set_status(status=trace.StatusCode.ERROR, description=err.__str__())
-        span.end()
+        self._detach_context_and_end_span(id_, span)
         return sp
 
 
@@ -324,6 +347,7 @@ def _start_otel(
         assert self.span_exporter is not None, (
             "span_exporter has to be non-null to be used within simple or batch span processors"
         )
+        span_processor: SpanProcessor
         if self.span_processor == "simple":
             span_processor = SimpleSpanProcessor(self.span_exporter)
         else:

llama-index-integrations/observability/llama-index-observability-otel/pyproject.toml

@@ -26,7 +26,7 @@ dev = [
 
 [project]
 name = "llama-index-observability-otel"
-version = "0.6.0"
+version = "0.6.1"
 description = "llama-index observability integration with OpenTelemetry"
 authors = [{name = "Clelia Astra Bertelli", email = "clelia@runllama.ai"}]
 requires-python = ">=3.10,<4.0"

llama-index-integrations/observability/llama-index-observability-otel/tests/test_otel.py

@@ -265,6 +265,44 @@ def make_handler():
     return exporter, handler, provider
 
 
+def enter_llama_span(
+    handler: OTelCompatibleSpanHandler,
+    span_id: str,
+    parent_id: Optional[str] = None,
+) -> None:
+    handler.span_enter(id_=span_id, bound_args=_bound, parent_id=parent_id)
+
+
+def exit_llama_span(handler: OTelCompatibleSpanHandler, span_id: str) -> None:
+    handler.span_exit(id_=span_id, bound_args=_bound)
+
+
+def finished_spans_by_name(
+    exporter: InMemorySpanExporter,
+    provider: TracerProvider,
+) -> dict[str, ReadableSpan]:
+    provider.force_flush()
+    spans = exporter.get_finished_spans()
+    spans_by_name = {span.name: span for span in spans}
+    assert len(spans_by_name) == len(spans)
+    return spans_by_name
+
+
+def assert_parent(child: ReadableSpan, parent: ReadableSpan) -> None:
+    assert child.parent is not None
+    assert child.parent.span_id == parent.context.span_id
+    assert child.parent.trace_id == parent.context.trace_id
+
+
+def assert_trace_chain(
+    spans_by_name: dict[str, ReadableSpan],
+    expected_chain: Sequence[str],
+) -> None:
+    assert set(spans_by_name) == set(expected_chain)
+    for parent_name, child_name in zip(expected_chain, expected_chain[1:]):
+        assert_parent(spans_by_name[child_name], spans_by_name[parent_name])
+
+
 def test_span_name_strips_uuid() -> None:
     exporter, handler, provider = make_handler()
     handler.span_enter(id_="MyWorkflow.run-abc123-def", bound_args=_bound)
@@ -357,17 +395,95 @@ def test_tags_not_mutated_by_new_span() -> None:
     assert tags == original_tags
 
 
+def test_span_enter_makes_otel_span_current_for_downstream_spans() -> None:
+    # Expected trace:
+    # root
+    #   downstream-child
+    exporter, handler, provider = make_handler()
+    clean_token = context.attach(context.Context())
+    try:
+        downstream_tracer = provider.get_tracer("downstream")
+
+        enter_llama_span(handler, "root-uuid")
+        root_otel_span = handler.all_spans["root-uuid"]
+        assert trace.get_current_span() is root_otel_span
+
+        with downstream_tracer.start_as_current_span("downstream-child"):
+            pass
+
+        exit_llama_span(handler, "root-uuid")
+
+        spans = finished_spans_by_name(exporter, provider)
+        assert_trace_chain(spans, ["root", "downstream-child"])
+        assert trace.get_current_span() is not root_otel_span
+    finally:
+        context.detach(clean_token)
+
+
+def test_otel_context_is_source_of_truth_when_external_spans_interleave() -> None:
+    # Expected trace:
+    # external-root
+    #   llama_parent
+    #     external-child
+    #       llama_child
+    exporter, handler, provider = make_handler()
+    clean_token = context.attach(context.Context())
+    tracer = provider.get_tracer("external")
+
+    try:
+        with tracer.start_as_current_span("external-root"):
+            enter_llama_span(handler, "llama_parent")
+            with tracer.start_as_current_span("external-child"):
+                enter_llama_span(
+                    handler,
+                    "llama_child",
+                    parent_id="llama_parent",
+                )
+                exit_llama_span(handler, "llama_child")
+            exit_llama_span(handler, "llama_parent")
+
+        spans = finished_spans_by_name(exporter, provider)
+        assert_trace_chain(
+            spans,
+            ["external-root", "llama_parent", "external-child", "llama_child"],
+        )
+    finally:
+        context.detach(clean_token)
+
+
+def test_span_exit_ends_span_when_context_detach_fails(monkeypatch: Any) -> None:
+    exporter, handler, provider = make_handler()
+    clean_token = context.attach(context.Context())
+    original_detach = context.detach
+
+    try:
+        handler.span_enter(id_="root-uuid", bound_args=_bound)
+
+        def fail_detach(token: Any) -> None:
+            raise RuntimeError("detach failed")
+
+        monkeypatch.setattr(context, "detach", fail_detach)
+        handler.span_exit(id_="root-uuid", bound_args=_bound)
+        provider.force_flush()
+
+        spans = exporter.get_finished_spans()
+        assert len(spans) == 1
+        assert spans[0].end_time is not None
+        assert handler.all_spans == {}
+        assert handler._context_tokens == {}
+    finally:
+        monkeypatch.setattr(context, "detach", original_detach)
+        context.detach(clean_token)
+
+
 def test_capture_propagation_context() -> None:
     """capture_propagation_context returns a dict with traceparent when a span is active."""
     exporter, handler, provider = make_handler()
     # Create a span so there's an active trace context
     handler.span_enter(id_="root-uuid", bound_args=_bound)
-    # Activate the OTel span in the current context so capture can see it
-    from opentelemetry.trace import set_span_in_context
 
     otel_span = handler.all_spans["root-uuid"]
-    ctx = set_span_in_context(otel_span)
-    context.attach(ctx)
+    assert trace.get_current_span() is otel_span
 
     captured = handler.capture_propagation_context()
     assert "otel" in captured
@@ -389,19 +505,15 @@ def test_capture_restore_propagation_roundtrip() -> None:
     - externally-set OTel context (e.g. baggage-like ambient values) propagates
       through the traceparent mechanism
     """
-    from opentelemetry.trace import set_span_in_context
-
     # --- Process A: create root span with tags, capture context ---
     exporter_a, handler_a, provider_a = make_handler()
 
     tags_a = {"handler_id": "h1", "run_id": "r1", "myapp.custom": "val_a"}
     original_tags_a = dict(tags_a)
     handler_a.span_enter(id_="root-uuid", bound_args=_bound, tags=tags_a)
 
-    # Activate the OTel span in ambient context (simulating what the Dispatcher
-    # would do before a serialization boundary)
     root_otel_span = handler_a.all_spans["root-uuid"]
-    context.attach(set_span_in_context(root_otel_span))
+    assert trace.get_current_span() is root_otel_span
 
     # Capture propagation context — this is what gets serialized across the boundary
     captured_ctx = handler_a.capture_propagation_context()
@@ -486,7 +598,6 @@ def test_dispatcher_propagation_roundtrip_with_tags() -> None:
         active_instrument_tags,
         instrument_tags,
     )
-    from opentelemetry.trace import set_span_in_context
 
     exporter_a, handler_a, provider_a = make_handler()
     exporter_b, handler_b, provider_b = make_handler()
@@ -508,9 +619,8 @@ def test_dispatcher_propagation_roundtrip_with_tags() -> None:
             id_="root-uuid", bound_args=_bound, tags=active_instrument_tags.get()
         )
 
-        # Activate OTel span in ambient context
         root_otel_span = handler_a.all_spans["root-uuid"]
-        context.attach(set_span_in_context(root_otel_span))
+        assert trace.get_current_span() is root_otel_span
 
         captured = dispatcher_a.capture_propagation_context()