solve tool output in stram, wrong total time

This commit is contained in:
Mery-Sanz 2025-04-30 10:13:41 +02:00
parent 695d49e4ed
commit 0cfa72bcf5
2 changed files with 51 additions and 21 deletions

View File

@ -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)

View File

@ -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