diff --git a/llama-index-core/llama_index/core/instrumentation/__init__.py b/llama-index-core/llama_index/core/instrumentation/__init__.py index a060bf2f095cbe089d5b852bbdc6887ee69c0268..b7564a02f7cd4eaec28a20ad2513bdfcdda072ea 100644 --- a/llama-index-core/llama_index/core/instrumentation/__init__.py +++ b/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, ) diff --git a/llama-index-core/llama_index/core/instrumentation/span_handlers/base.py b/llama-index-core/llama_index/core/instrumentation/span_handlers/base.py index f18cf658adeaa9466eeb8ba3d13a9de75dc942df..a54a8d7b30b79906ddcf4834718d0368ff924a68 100644 --- a/llama-index-core/llama_index/core/instrumentation/span_handlers/base.py +++ b/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.""" ... diff --git a/llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py b/llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py index a1ceafd879916855619639832ce66af4aa68a8fe..96c54b2fb2982f8389b5b0294d9f3868ffee7358 100644 --- a/llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py +++ b/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."""