rename and refactor according to review tips

packages/ragbits-core/src/ragbits/core/audit/__init__.py

@@ -44,9 +44,9 @@ def set_trace_handlers(handlers: Handler | list[Handler]) -> None:
                 _trace_handlers.append(OtelTraceHandler())
 
             elif handler == "cli":
-                from ragbits.core.audit.cli_tracer import CLITracer
+                from ragbits.core.audit.cli import CLITraceHandler
 
-                _trace_handlers.append(CLITracer())
+                _trace_handlers.append(CLITraceHandler())
 
             else:
                 raise ValueError(f"Handler {handler} not found.")

packages/ragbits-core/src/ragbits/core/audit/cli_span.py → packages/ragbits-core/src/ragbits/core/audit/cli.py

@@ -1,8 +1,11 @@
 import time
 from typing import Optional
 
+from rich import print as rich_print
 from rich.tree import Tree
 
+from ragbits.core.audit import TraceHandler
+
 
 class CLISpan:
     """
@@ -22,7 +25,7 @@ def __init__(self, name: str, inputs: dict, parent: Optional["CLISpan"] = None):
         """
         self.name = name
         self.parent = parent
-        self.start_time: float = time.time()
+        self.start_time: float = time.perf_counter()
         self.end_time: float | None = None
         self.children: list[CLISpan] = []
         self.status: str = "started"
@@ -36,7 +39,7 @@ def end(self) -> None:
         further calls are ignored.
         """
         if self.end_time is None:
-            self.end_time = time.time()
+            self.end_time = time.perf_counter()
 
     def to_tree(self, tree: Tree | None = None, color: str = "bold blue") -> Tree | None:
         """
@@ -53,16 +56,17 @@ def to_tree(self, tree: Tree | None = None, color: str = "bold blue") -> Tree |
         secondary_color = "grey50"
         error_color = "bold red"
         child_color = "bold green"
+        duration = self.end_time - self.start_time if self.end_time else 0.0
 
-        if tree is None and self.end_time:
+        if tree is None:
             tree = Tree(
-                f"[{color}]{self.name}[/{color}] Duration: {self.end_time - self.start_time:.3f}s\n"
+                f"[{color}]{self.name}[/{color}] Duration: {duration:.3f}s\n"
                 f"[{secondary_color}]Inputs: {self.inputs}\nOutputs: {self.outputs})[/{secondary_color}]"
             )
 
-        elif tree and self.end_time:
+        else:
             child_tree = tree.add(
-                f"[{color}]{self.name}[/{color}] Duration: {self.end_time - self.start_time:.3f}s\n"
+                f"[{color}]{self.name}[/{color}] Duration: {duration:.3f}s\n"
                 f"[{secondary_color}]Inputs: {self.inputs}\nOutputs: {self.outputs})[/{secondary_color}]"
             )
             tree = child_tree
@@ -73,3 +77,56 @@ def to_tree(self, tree: Tree | None = None, color: str = "bold blue") -> Tree |
             else:
                 child.to_tree(tree, child_color)
         return tree
+
+
+class CLITraceHandler(TraceHandler[CLISpan]):
+    """
+    CLITraceHandler class for all trace handlers.
+    """
+
+    def start(self, name: str, inputs: dict, current_span: CLISpan | None = None) -> CLISpan:  # noqa: PLR6301
+        """
+        Log input data at the beginning of the trace.
+
+        Args:
+            name: The name of the trace.
+            inputs: The input data.
+            current_span: The current trace span.
+
+        Returns:
+            The updated current trace span.
+        """
+        span = CLISpan(name, inputs, current_span)
+
+        if current_span:
+            current_span.children.append(span)
+
+        return span
+
+    def stop(self, outputs: dict, current_span: CLISpan) -> None:  # noqa: PLR6301
+        """
+        Log output data at the end of the trace.
+
+        Args:
+            outputs: The output data.
+            current_span: The current trace span.
+        """
+        current_span.end()
+        current_span.status = "done"
+        current_span.outputs = outputs
+
+        if current_span.parent is None:
+            rich_print(current_span.to_tree())
+
+    def error(self, error: Exception, current_span: CLISpan) -> None:  # noqa: PLR6301
+        """
+        Log error during the trace.
+
+        Args:
+            error: The error that occurred.
+            current_span: The current trace span.
+        """
+        current_span.end()
+        current_span.status = "error"
+        if current_span.parent is None:
+            rich_print(current_span.to_tree())