From dce35afdd34600d991add53fccaf0a20e299d86f Mon Sep 17 00:00:00 2001
From: Roger Yang <roger.yang@arize.com>
Date: Thu, 12 Dec 2024 08:44:41 -0800
Subject: [PATCH 1/2] fix: get_current_span when not instrumented

---
 .../instrumentation/llama_index/__init__.py    |  7 ++++++-
 .../llama_index/test_get_current_span.py       | 18 +++++-------------
 2 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/__init__.py b/python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/__init__.py
index d98d05251..372501558 100644
--- a/python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/__init__.py
+++ b/python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/__init__.py
@@ -109,9 +109,14 @@ def _uninstrument(self, **kwargs: Any) -> None:
 
 def get_current_span() -> Optional[Span]:
     from llama_index.core.instrumentation.span import active_span_id
+    from openinference.instrumentation.llama_index._handler import _SpanHandler
 
     if not isinstance(id_ := active_span_id.get(), str):
         return None
-    if (span := LlamaIndexInstrumentor()._span_handler.open_spans.get(id_)) is None:
+    instrumentor = LlamaIndexInstrumentor()
+    span_handler = getattr(instrumentor, "_span_handler", None)
+    if not isinstance(span_handler, _SpanHandler):
+        return None
+    if (span := span_handler.open_spans.get(id_)) is None:
         return None
     return span._otel_span
diff --git a/python/instrumentation/openinference-instrumentation-llama-index/tests/openinference/instrumentation/llama_index/test_get_current_span.py b/python/instrumentation/openinference-instrumentation-llama-index/tests/openinference/instrumentation/llama_index/test_get_current_span.py
index 87ddce3a8..57bbb4553 100644
--- a/python/instrumentation/openinference-instrumentation-llama-index/tests/openinference/instrumentation/llama_index/test_get_current_span.py
+++ b/python/instrumentation/openinference-instrumentation-llama-index/tests/openinference/instrumentation/llama_index/test_get_current_span.py
@@ -1,8 +1,6 @@
 from asyncio import create_task, 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
@@ -14,7 +12,7 @@
 
 
 @dispatcher.span  # type: ignore[misc,unused-ignore]
-async def foo(k: int) -> str:
+async def foo(k: int = 1) -> str:
     child = create_task(foo(k - 1)) if k > 1 else None
     await sleep(random() / 100)
     span = get_current_span()
@@ -24,10 +22,14 @@ async def foo(k: int) -> str:
 
 
 async def test_get_current_span(
+    tracer_provider: TracerProvider,
     in_memory_span_exporter: InMemorySpanExporter,
 ) -> None:
+    assert await foo() == ""
     n, k = 10, 5
+    LlamaIndexInstrumentor().instrument(tracer_provider=tracer_provider)
     await gather(*(foo(k) for _ in range(n)))
+    LlamaIndexInstrumentor().uninstrument()
     spans = in_memory_span_exporter.get_finished_spans()
     assert len(spans) == n * k
     seen = set()
@@ -38,14 +40,4 @@ async def test_get_current_span(
         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()
-
-
 OUTPUT_VALUE = SpanAttributes.OUTPUT_VALUE

From b71ee4bebbf84102798ad4f447527654e061bde1 Mon Sep 17 00:00:00 2001
From: Roger Yang <roger.yang@arize.com>
Date: Thu, 12 Dec 2024 08:48:10 -0800
Subject: [PATCH 2/2] clean up

---
 .../openinference/instrumentation/llama_index/__init__.py    | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/__init__.py b/python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/__init__.py
index 372501558..50c09c9b6 100644
--- a/python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/__init__.py
+++ b/python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/__init__.py
@@ -114,7 +114,10 @@ def get_current_span() -> Optional[Span]:
     if not isinstance(id_ := active_span_id.get(), str):
         return None
     instrumentor = LlamaIndexInstrumentor()
-    span_handler = getattr(instrumentor, "_span_handler", None)
+    try:
+        span_handler = instrumentor._span_handler
+    except AttributeError:
+        return None
     if not isinstance(span_handler, _SpanHandler):
         return None
     if (span := span_handler.open_spans.get(id_)) is None: