diff --git a/packages/ragbits-core/src/ragbits/core/audit/__init__.py b/packages/ragbits-core/src/ragbits/core/audit/__init__.py index 24bc2b517..5177426ff 100644 --- a/packages/ragbits-core/src/ragbits/core/audit/__init__.py +++ b/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.") diff --git a/packages/ragbits-core/src/ragbits/core/audit/cli_span.py b/packages/ragbits-core/src/ragbits/core/audit/cli.py similarity index 54% rename from packages/ragbits-core/src/ragbits/core/audit/cli_span.py rename to packages/ragbits-core/src/ragbits/core/audit/cli.py index 9992b4ea9..cda35790b 100644 --- a/packages/ragbits-core/src/ragbits/core/audit/cli_span.py +++ b/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()) diff --git a/packages/ragbits-core/src/ragbits/core/audit/cli_tracer.py b/packages/ragbits-core/src/ragbits/core/audit/cli_tracer.py deleted file mode 100644 index dff67f04b..000000000 --- a/packages/ragbits-core/src/ragbits/core/audit/cli_tracer.py +++ /dev/null @@ -1,66 +0,0 @@ -from rich.console import Console - -from ragbits.core.audit import TraceHandler -from ragbits.core.audit.cli_span import CLISpan - - -class CLITracer(TraceHandler[CLISpan]): - """ - CLITracer class for all trace handlers. - """ - - def __init__(self) -> None: - """ - Construct a new CLITracer instance. - """ - super().__init__() - - def start(self, name: str, inputs: dict, current_span: CLISpan | None = None) -> CLISpan: - """ - 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. - """ - parent = self._spans.get()[-1] if self._spans.get() else None - span = CLISpan(name, inputs, parent) - - if parent: - parent.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: - console = Console() - console.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: - console = Console() - console.print(current_span.to_tree())