From 63f6ecc9575e232791b33ca7b9d0fd96943a76c2 Mon Sep 17 00:00:00 2001
From: Roger Yang <roger.yang@arize.com>
Date: Fri, 19 Jan 2024 11:01:40 -0800
Subject: [PATCH] clean up

---
 .../pyproject.toml                                          | 3 +--
 .../openinference/instrumentation/llama_index/_callback.py  | 6 ++----
 .../instrumentation/llama_index/test_callback.py            | 4 +---
 3 files changed, 4 insertions(+), 9 deletions(-)

diff --git a/python/instrumentation/openinference-instrumentation-llama-index/pyproject.toml b/python/instrumentation/openinference-instrumentation-llama-index/pyproject.toml
index b4ef07ee4..c54b098ea 100644
--- a/python/instrumentation/openinference-instrumentation-llama-index/pyproject.toml
+++ b/python/instrumentation/openinference-instrumentation-llama-index/pyproject.toml
@@ -38,9 +38,8 @@ instruments = [
 test = [
   "llama-index == 0.9.14",
   "openai",
-  "openinference-instrumentation-openai",
   "opentelemetry-sdk",
-  "opentelemetry-instrumentation-httpx",
+  "openinference-instrumentation-openai",
   "respx",
 ]
 
diff --git a/python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/_callback.py b/python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/_callback.py
index d5d4a74d8..f28ee96b1 100644
--- a/python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/_callback.py
+++ b/python/instrumentation/openinference-instrumentation-llama-index/src/openinference/instrumentation/llama_index/_callback.py
@@ -304,15 +304,13 @@ def start_trace(self, trace_id: Optional[str] = None) -> None:
                 self._stragglers.pop(event_id)
 
     def end_trace(self, trace_id: Optional[str] = None, *args: Any, **kwargs: Any) -> None:
-        dfs_stack, adjacency_list = _build_graph(self._event_data)
+        roots, adjacency_list = _build_graph(self._event_data)
         self._event_data.clear()
+        dfs_stack = roots.copy()
         while dfs_stack:
             event_id, event_data = dfs_stack.pop()
             event_type = event_data.event_type
             attributes = event_data.attributes
-            if not event_data.span:
-                # `on_event_start` was never called, so discard the event.
-                continue
 
             if event_type is CBEventType.LLM:
                 while dfs_stack:
diff --git a/python/instrumentation/openinference-instrumentation-llama-index/tests/openinference/instrumentation/llama_index/test_callback.py b/python/instrumentation/openinference-instrumentation-llama-index/tests/openinference/instrumentation/llama_index/test_callback.py
index 48510a480..7f85b4c43 100644
--- a/python/instrumentation/openinference-instrumentation-llama-index/tests/openinference/instrumentation/llama_index/test_callback.py
+++ b/python/instrumentation/openinference-instrumentation-llama-index/tests/openinference/instrumentation/llama_index/test_callback.py
@@ -10,9 +10,7 @@
 from httpx import AsyncByteStream, Response, SyncByteStream
 from llama_index import ListIndex, ServiceContext
 from llama_index.callbacks import CallbackManager
-from llama_index.llms import (
-    OpenAI,
-)
+from llama_index.llms import OpenAI
 from llama_index.schema import Document, TextNode
 from openinference.instrumentation.llama_index import LlamaIndexInstrumentor
 from openinference.instrumentation.openai import OpenAIInstrumentor