feat: add get_current_span helper function for llama-index

python/instrumentation/openinference-instrumentation-llama-index/pyproject.toml

@@ -66,6 +66,7 @@ packages = ["src/openinference"]
 
 [tool.pytest.ini_options]
 asyncio_mode = "auto"
+asyncio_default_fixture_loop_scope = "function"
 testpaths = [
   "tests",
 ]

python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/__init__.py

@@ -1,8 +1,9 @@
 import logging
-from typing import Any, Collection
+from typing import Any, Collection, Optional
 
 from opentelemetry import trace as trace_api
 from opentelemetry.instrumentation.instrumentor import BaseInstrumentor  # type: ignore
+from opentelemetry.trace import Span
 
 from openinference.instrumentation import OITracer, TraceConfig
 from openinference.instrumentation.llama_index.package import _instruments
@@ -104,3 +105,13 @@ def _uninstrument(self, **kwargs: Any) -> None:
                 dispatcher.event_handlers,
             )
             self._event_handler = None
+
+
+def get_current_span() -> Optional[Span]:
+    from llama_index.core.instrumentation.span import active_span_id
+
+    if not isinstance(id_ := active_span_id.get(), str):
+        return None
+    if (span := LlamaIndexInstrumentor()._span_handler.open_spans.get(id_)) is None:
+        return None
+    return span._otel_span

python/instrumentation/openinference-instrumentation-llama-index/tests/openinference/instrumentation/llama_index/conftest.py

@@ -0,0 +1,18 @@
+import pytest
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export import SimpleSpanProcessor
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+
+@pytest.fixture
+def in_memory_span_exporter() -> InMemorySpanExporter:
+    return InMemorySpanExporter()
+
+
+@pytest.fixture
+def tracer_provider(
+    in_memory_span_exporter: InMemorySpanExporter,
+) -> TracerProvider:
+    tracer_provider = TracerProvider()
+    tracer_provider.add_span_processor(SimpleSpanProcessor(in_memory_span_exporter))
+    return tracer_provider

python/instrumentation/openinference-instrumentation-llama-index/tests/openinference/instrumentation/llama_index/test_get_current_span.py

@@ -0,0 +1,45 @@
+from asyncio import gather, sleep
+from random import random
+from typing import Iterator
+
+import pytest
+from llama_index.core.instrumentation import get_dispatcher
+from opentelemetry.sdk.trace import TracerProvider
+from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter
+
+from openinference.instrumentation.llama_index import LlamaIndexInstrumentor, get_current_span
+
+dispatcher = get_dispatcher(__name__)
+
+
+@dispatcher.span  # type: ignore[misc,unused-ignore]
+async def foo(k: int) -> str:
+    if k > 1:
+        await gather(sleep(random() / 100), foo(k - 1))
+    if (span := get_current_span()) is None:
+        return ""
+    return str(span.get_span_context().span_id)
+
+
+async def test_get_current_span(
+    in_memory_span_exporter: InMemorySpanExporter,
+) -> None:
+    n, k = 10, 10
+    await gather(*(foo(k) for _ in range(n)))
+    assert len(spans := in_memory_span_exporter.get_finished_spans()) == n * k
+    seen = set()
+    for span in spans:
+        assert span.attributes and span.context
+        assert (expected := str(span.context.span_id)) not in seen
+        seen.add(expected)
+        assert span.attributes.get("output.value") == expected
+
+
+@pytest.fixture(autouse=True)
+def instrument(
+    tracer_provider: TracerProvider,
+    in_memory_span_exporter: InMemorySpanExporter,
+) -> Iterator[None]:
+    LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)
+    yield
+    LlamaIndexInstrumentor().uninstrument()