Skip to content

Span handlers

Bases: BaseModel, Generic[T]

Source code in llama-index-core/llama_index/core/instrumentation/span_handlers/base.py 
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
class BaseSpanHandler(BaseModel, Generic[T]):
    open_spans: Dict[str, T] = Field(
        default_factory=dict, description="Dictionary of open spans."
    )
    current_span_id: Optional[str] = Field(
        default=None, description="Id of current span."
    )

    class Config:
        arbitrary_types_allowed = True

    def class_name(cls) -> str:
        """Class name."""
        return "BaseSpanHandler"

    def span_enter(self, *args, id: str, **kwargs) -> None:
        """Logic for entering a span."""
        if id in self.open_spans:
            pass  # should probably raise an error here
        else:
            # TODO: thread safe?
            span = self.new_span(
                *args, id=id, parent_span_id=self.current_span_id, **kwargs
            )
            if span:
                self.open_spans[id] = span
                self.current_span_id = id

    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]

    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]

    @abstractmethod
    def new_span(
        self, *args, id: str, parent_span_id: Optional[str], **kwargs
    ) -> Optional[T]:
        """Create a span."""
        ...

    @abstractmethod
    def prepare_to_exit_span(
        self, *args, id: str, result: Optional[Any] = None, **kwargs
    ) -> Any:
        """Logic for preparing to exit a span."""
        ...

    @abstractmethod
    def prepare_to_drop_span(
        self, *args, id: str, err: Optional[Exception], **kwargs
    ) -> Any:
        """Logic for preparing to drop a span."""
        ...

class_name #

class_name() -> str

Class name.

Source code in llama-index-core/llama_index/core/instrumentation/span_handlers/base.py 
21
22
23
def class_name(cls) -> str:
    """Class name."""
    return "BaseSpanHandler"

span_enter #

span_enter(*args, id: str, **kwargs) -> None

Logic for entering a span.

Source code in llama-index-core/llama_index/core/instrumentation/span_handlers/base.py 
25
26
27
28
29
30
31
32
33
34
35
36
def span_enter(self, *args, id: str, **kwargs) -> None:
    """Logic for entering a span."""
    if id in self.open_spans:
        pass  # should probably raise an error here
    else:
        # TODO: thread safe?
        span = self.new_span(
            *args, id=id, parent_span_id=self.current_span_id, **kwargs
        )
        if span:
            self.open_spans[id] = span
            self.current_span_id = id

span_exit #

span_exit(*args, id: str, result: Optional[Any] = None, **kwargs) -> None

Logic for exiting a span.

Source code in llama-index-core/llama_index/core/instrumentation/span_handlers/base.py 
38
39
40
41
42
43
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_drop #

span_drop(*args, id: str, err: Optional[Exception], **kwargs) -> None

Logic for dropping a span i.e. early exit.

Source code in llama-index-core/llama_index/core/instrumentation/span_handlers/base.py 
45
46
47
48
49
50
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]

new_span abstractmethod #

new_span(*args, id: str, parent_span_id: Optional[str], **kwargs) -> Optional[T]

Create a span.

Source code in llama-index-core/llama_index/core/instrumentation/span_handlers/base.py 
52
53
54
55
56
57
@abstractmethod
def new_span(
    self, *args, id: str, parent_span_id: Optional[str], **kwargs
) -> Optional[T]:
    """Create a span."""
    ...

prepare_to_exit_span abstractmethod #

prepare_to_exit_span(*args, id: str, result: Optional[Any] = None, **kwargs) -> Any

Logic for preparing to exit a span.

Source code in llama-index-core/llama_index/core/instrumentation/span_handlers/base.py 
59
60
61
62
63
64
@abstractmethod
def prepare_to_exit_span(
    self, *args, id: str, result: Optional[Any] = None, **kwargs
) -> Any:
    """Logic for preparing to exit a span."""
    ...

prepare_to_drop_span abstractmethod #

prepare_to_drop_span(*args, id: str, err: Optional[Exception], **kwargs) -> Any

Logic for preparing to drop a span.

Source code in llama-index-core/llama_index/core/instrumentation/span_handlers/base.py 
66
67
68
69
70
71
@abstractmethod
def prepare_to_drop_span(
    self, *args, id: str, err: Optional[Exception], **kwargs
) -> Any:
    """Logic for preparing to drop a span."""
    ...

Bases: BaseSpanHandler[SimpleSpan]

Span Handler that managest SimpleSpan's.

Source code in llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py 
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
class SimpleSpanHandler(BaseSpanHandler[SimpleSpan]):
    """Span Handler that managest SimpleSpan's."""

    completed_spans: List[SimpleSpan] = Field(
        default_factory=list, description="List of completed spans."
    )

    def class_name(cls) -> str:
        """Class name."""
        return "SimpleSpanHandler"

    def new_span(
        self, *args, id: str, parent_span_id: Optional[str], **kwargs
    ) -> SimpleSpan:
        """Create a span."""
        return SimpleSpan(id_=id, parent_id=parent_span_id)

    def prepare_to_exit_span(
        self, *args, id: str, result: Optional[Any] = None, **kwargs
    ) -> None:
        """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]

    def prepare_to_drop_span(
        self, *args, id: str, err: Optional[Exception], **kwargs
    ) -> None:
        """Logic for droppping a span."""
        if err:
            raise err

    def _get_trace_trees(self) -> List["Tree"]:
        """Method for getting trace trees."""
        try:
            from treelib import Tree
            from treelib.exceptions import NodeIDAbsentError
        except ImportError as e:
            raise ImportError(
                "`treelib` package is missing. Please install it by using "
                "`pip install treelib`."
            )
        sorted_spans = sorted(self.completed_spans, key=lambda x: x.start_time)

        trees = []
        tree = Tree()
        for span in sorted_spans:
            if span.parent_id is None:
                # complete old tree unless its empty (i.e., start of loop)
                if tree.all_nodes():
                    trees.append(tree)
                    # start new tree
                    tree = Tree()

            try:
                tree.create_node(
                    tag=f"{span.id_} ({span.duration})",
                    identifier=span.id_,
                    parent=span.parent_id,
                    data=span.start_time,
                )
            except NodeIDAbsentError:
                warnings.warn("Parent with id {span.parent_id} missing from spans")
                # create new tree and fake parent node
                trees.append(tree)
                tree = Tree()
                tree.create_node(
                    tag=f"{span.parent_id} (MISSING)",
                    identifier=span.parent_id,
                    parent=None,
                    data=span.start_time,
                )
                tree.create_node(
                    tag=f"{span.id_} ({span.duration})",
                    identifier=span.id_,
                    parent=span.parent_id,
                    data=span.start_time,
                )

        trees.append(tree)
        return trees

    def print_trace_trees(self) -> None:
        """Method for viewing trace trees."""
        trees = self._get_trace_trees()
        for tree in trees:
            print(tree.show(stdout=False, sorting=True, key=lambda node: node.data))
            print("")

class_name #

class_name() -> str

Class name.

Source code in llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py 
19
20
21
def class_name(cls) -> str:
    """Class name."""
    return "SimpleSpanHandler"

new_span #

new_span(*args, id: str, parent_span_id: Optional[str], **kwargs) -> SimpleSpan

Create a span.

Source code in llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py 
23
24
25
26
27
def new_span(
    self, *args, id: str, parent_span_id: Optional[str], **kwargs
) -> SimpleSpan:
    """Create a span."""
    return SimpleSpan(id_=id, parent_id=parent_span_id)

prepare_to_exit_span #

prepare_to_exit_span(*args, id: str, result: Optional[Any] = None, **kwargs) -> None

Logic for preparing to drop a span.

Source code in llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py 
29
30
31
32
33
34
35
36
37
def prepare_to_exit_span(
    self, *args, id: str, result: Optional[Any] = None, **kwargs
) -> None:
    """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]

prepare_to_drop_span #

prepare_to_drop_span(*args, id: str, err: Optional[Exception], **kwargs) -> None

Logic for droppping a span.

Source code in llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py 
39
40
41
42
43
44
def prepare_to_drop_span(
    self, *args, id: str, err: Optional[Exception], **kwargs
) -> None:
    """Logic for droppping a span."""
    if err:
        raise err

print_trace_trees #

print_trace_trees() -> None

Method for viewing trace trees.

Source code in llama-index-core/llama_index/core/instrumentation/span_handlers/simple.py 
 96
 97
 98
 99
100
101
def print_trace_trees(self) -> None:
    """Method for viewing trace trees."""
    trees = self._get_trace_trees()
    for tree in trees:
        print(tree.show(stdout=False, sorting=True, key=lambda node: node.data))
        print("")