clean up

Arize-ai · Jan 19, 2024 · 63f6ecc · 63f6ecc
1 parent 322a3cb
commit 63f6ecc
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 9 deletions.
diff --git 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/...ce-instrumentation-llama-index/src/openinference/instrumentation/llama_index/_callback.py b/...ce-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/...trumentation-llama-index/tests/openinference/instrumentation/llama_index/test_callback.py b/...trumentation-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