Small fixes to instrumentation (#12287)

span_handlers in root dispatcher; return Optional[T] to handle null handler case

llama-index-core/llama_index/core/instrumentation/__init__.py

@@ -5,7 +5,7 @@ from llama_index.core.instrumentation.span_handlers import NullSpanHandler
 root_dispatcher: Dispatcher = Dispatcher(
     name="root",
     event_handlers=[NullEventHandler()],
-    span_handler=NullSpanHandler(),
+    span_handlers=[NullSpanHandler()],
     propagate=False,
 )
 

llama-index-core/llama_index/core/instrumentation/span_handlers/base.py

@@ -37,17 +37,19 @@ class BaseSpanHandler(BaseModel, Generic[T]):
 
     def span_exit(self, *args, id: str, result: Optional[Any] = None, **kwargs) -> None:
         """Logic for exiting a span."""
-        self.prepare_to_exit_span(*args, id=id, result=result, **kwargs)
-        if self.current_span_id == id:
-            self.current_span_id = self.open_spans[id].parent_id
-        del self.open_spans[id]
+        span = self.prepare_to_exit_span(*args, id=id, result=result, **kwargs)
+        if span:
+            if self.current_span_id == id:
+                self.current_span_id = self.open_spans[id].parent_id
+            del self.open_spans[id]
 
     def span_drop(self, *args, id: str, err: Optional[Exception], **kwargs) -> None:
         """Logic for dropping a span i.e. early exit."""
-        self.prepare_to_drop_span(*args, id=id, err=err, **kwargs)
-        if self.current_span_id == id:
-            self.current_span_id = self.open_spans[id].parent_id
-        del self.open_spans[id]
+        span = self.prepare_to_drop_span(*args, id=id, err=err, **kwargs)
+        if span:
+            if self.current_span_id == id:
+                self.current_span_id = self.open_spans[id].parent_id
+            del self.open_spans[id]
 
     @abstractmethod
     def new_span(
@@ -59,13 +61,13 @@ class BaseSpanHandler(BaseModel, Generic[T]):
     @abstractmethod
     def prepare_to_exit_span(
         self, *args, id: str, result: Optional[Any] = None, **kwargs
-    ) -> Any:
+    ) -> Optional[T]:
         """Logic for preparing to exit a span."""
         ...
 
     @abstractmethod
     def prepare_to_drop_span(
         self, *args, id: str, err: Optional[Exception], **kwargs
-    ) -> Any:
+    ) -> Optional[T]:
         """Logic for preparing to drop a span."""
         ...

llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py

@@ -28,20 +28,22 @@ class SimpleSpanHandler(BaseSpanHandler[SimpleSpan]):
 
     def prepare_to_exit_span(
         self, *args, id: str, result: Optional[Any] = None, **kwargs
-    ) -> None:
+    ) -> SimpleSpan:
         """Logic for preparing to drop a span."""
         span = self.open_spans[id]
         span = cast(SimpleSpan, span)
         span.end_time = datetime.now()
         span.duration = (span.end_time - span.start_time).total_seconds()
         self.completed_spans += [span]
+        return span
 
     def prepare_to_drop_span(
         self, *args, id: str, err: Optional[Exception], **kwargs
-    ) -> None:
+    ) -> SimpleSpan:
         """Logic for droppping a span."""
-        if err:
-            raise err
+        if id in self.open_spans:
+            return self.open_spans[id]
+        return None
 
     def _get_trace_trees(self) -> List["Tree"]:
         """Method for getting trace trees."""