diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index ba99dd04..8fd02bfc 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -12,6 +12,8 @@ import time import uuid import sys from wasabi import color # pylint: disable=import-error +from cai.util import format_time +from cai.cli import START_TIME # Global dictionary to store active sessions ACTIVE_SESSIONS = {} @@ -318,7 +320,17 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): header.append("(", style="yellow") header.append(args_str, style="yellow") header.append(")", style="yellow") - + start_time = time.time() + tool_time = 0 + total_time = START_TIME + timing_info = [] + if total_time: + timing_info.append(f"Total: {format_time(total_time)}") + if tool_time: + timing_info.append(f"Tool: {format_time(tool_time)}") + if timing_info: + header.append(f" [{' | '.join(timing_info)}]", style="cyan") + content = Text() panel = Panel( @@ -333,15 +345,35 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): # Start Live display with Live(panel, console=console, refresh_per_second=4) as live: # Stream stdout in real-time + start_time = time.time() for line in iter(process.stdout.readline, ''): if not line: break - + # Add to output collection output_buffer.append(line) - + # Update content with new line content.append(line, style="bright_white") + + # Update tool_time and header with new timing info + tool_time = time.time() - start_time + total_time = time.time() - start_time + tool_time + # Remove any previous timing info from header (rebuild header) + timing_info = [] + if total_time: + timing_info.append(f"Total: {format_time(total_time)}") + if tool_time: + timing_info.append(f"Tool: {format_time(tool_time)}") + # Rebuild header to update timing + header = Text() + header.append(tool_name, style="#00BCD4") + header.append("(", style="yellow") + header.append(args_str, style="yellow") + header.append(")", style="yellow") + if timing_info: + header.append(f" [{' | '.join(timing_info)}]", style="cyan") + panel = Panel( Text.assemble(header, "\n\n", content), title="[bold green]Tool Execution[/bold green]", @@ -351,7 +383,6 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): box=ROUNDED ) live.update(panel) - # Check if process is done process.stdout.close() return_code = process.wait(timeout=timeout) diff --git a/src/cai/util.py b/src/cai/util.py index 5a3bec9d..e7a64e1c 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -522,7 +522,22 @@ def get_model_name(model): return model # If not a string, use environment variable return os.environ.get('CAI_MODEL', 'qwen2.5:72b') - + # Helper function to format time in a human-readable way +def format_time(seconds): + if seconds is None: + return "N/A" + + if seconds < 60: + return f"{seconds:.1f}s" + elif seconds < 3600: + minutes = int(seconds / 60) + seconds_remainder = seconds % 60 + return f"{minutes}m {seconds_remainder:.1f}s" + else: + hours = int(seconds / 3600) + minutes = int((seconds % 3600) / 60) + return f"{hours}h {minutes}m" + def get_model_pricing(model_name): """ Get pricing information for a model, using the CostTracker's implementation. @@ -1161,22 +1176,6 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut else: args_str = str(args) - # Helper function to format time in a human-readable way - def format_time(seconds): - if seconds is None: - return "N/A" - - if seconds < 60: - return f"{seconds:.1f}s" - elif seconds < 3600: - minutes = int(seconds / 60) - seconds_remainder = seconds % 60 - return f"{minutes}m {seconds_remainder:.1f}s" - else: - hours = int(seconds / 3600) - minutes = int((seconds % 3600) / 60) - return f"{hours}h {minutes}m" - # Get session timing information try: from cai.cli import START_TIME