diff --git a/src/cai/cli.py b/src/cai/cli.py index 72940444..53c9ce05 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -100,6 +100,7 @@ Usage Examples: import os import sys +import time from dotenv import load_dotenv from openai import AsyncOpenAI from cai.sdk.agents import OpenAIChatCompletionsModel, Agent, Runner, AsyncOpenAI @@ -114,7 +115,7 @@ from cai.util import create_agent_streaming_context, update_agent_streaming_cont from cai.repl.commands import FuzzyCommandCompleter, handle_command as commands_handle_command from cai.repl.ui.keybindings import create_key_bindings from cai.repl.ui.logging import setup_session_logging -from cai.repl.ui.banner import display_banner +from cai.repl.ui.banner import display_banner, display_quick_guide from cai.repl.ui.prompt import get_user_input from cai.repl.ui.toolbar import get_toolbar_with_refresh @@ -132,6 +133,10 @@ load_dotenv() # # set_default_openai_client(external_client) +# Global variables for timing tracking +global START_TIME +START_TIME = time.time() + set_tracing_disabled(True) # llm_model=os.getenv('LLM_MODEL', 'gpt-4o-mini') @@ -167,7 +172,8 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= """ agent = starting_agent turn_count = 0 - + ACTIVE_TIME = 0 + idle_time = 0 console = Console() # Initialize command completer and key bindings @@ -180,7 +186,7 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= # Display banner display_banner(console) - + display_quick_guide(console) # Function to get the short name of the agent for display def get_agent_short_name(agent): if hasattr(agent, 'name'): @@ -201,6 +207,7 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= while turn_count < max_turns: try: + idle_start_time = time.time() # Get user input with command completion and history user_input = get_user_input( command_completer, @@ -209,8 +216,59 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= get_toolbar_with_refresh, current_text ) + idle_time += time.time() - idle_start_time + except KeyboardInterrupt: + def format_time(seconds): + mins, secs = divmod(int(seconds), 60) + hours, mins = divmod(mins, 60) + return f"{hours:02d}:{mins:02d}:{secs:02d}" + + Total = time.time() - START_TIME + idle_time += time.time() - idle_start_time + try: + active_time = Total - idle_time + + metrics = { + "session_time": format_time(Total), + "active_time": format_time(active_time), + "idle_time": format_time(idle_time), + "llm_time": "0.0s", # Placeholder, update if available + "llm_percentage": 0.0, # Placeholder, update if available + } + logging_path = None # Set this if you have a log file path + + content = [] + content.append(f"Session Time: {metrics['session_time']}") + content.append(f"Active Time: {metrics['active_time']}") + content.append(f"Idle Time: {metrics['idle_time']}") + if logging_path: + content.append(f"Log available at: {logging_path}") + + def print_session_summary(console, metrics, logging_path=None): + """ + Print a session summary panel using Rich. + """ + from rich.panel import Panel + from rich.text import Text + from rich.box import ROUNDED + from rich.console import Group + + time_panel = Panel( + Group(*[Text(line) for line in content]), + border_style="blue", + box=ROUNDED, + padding=(0, 1), + title="[bold]Session Summary[/bold]", + title_align="left" + ) + console.print(time_panel) + + print_session_summary(console, metrics, logging_path) + except Exception: + pass break + try: # Handle special commands if user_input.startswith('/') or user_input.startswith('$'): @@ -355,10 +413,10 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= asyncio.run(process_streamed_response()) else: # Use non-streamed response - console.print("[dim]Thinking...[/dim]") response = asyncio.run(Runner.run(agent, user_input)) #console.print(f"Agent: {response.final_output}") # NOTE: this line is commented to avoid duplicate output turn_count += 1 + except KeyboardInterrupt: if stream: # Ensure streaming context is cleaned up on keyboard interrupt diff --git a/src/cai/repl/ui/banner.py b/src/cai/repl/ui/banner.py index 006f83cd..b750c4cb 100644 --- a/src/cai/repl/ui/banner.py +++ b/src/cai/repl/ui/banner.py @@ -253,3 +253,107 @@ def display_welcome_tips(console: Console): title="Quick Tips", border_style="blue" )) + + +def display_quick_guide(console: Console): + """Display the quick guide.""" + # Display help panel instead + from rich.panel import Panel + from rich.text import Text + from rich.columns import Columns + from rich.console import Group # <-- Fix: import Group + + help_text = Text.assemble( + ("CAI Command Reference", "bold cyan underline"), "\n\n", + ("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━", "dim"), "\n", + ("WORKSPACE", "bold yellow"), "\n", + (" CAI>/ws set [NAME]", "green"), " - Set current workspace directory\n\n", + ("AGENT MANAGEMENT", "bold yellow"), "\n", + (" CAI>/agent [NAME]", "green"), " - Switch to specific agent by name\n", + (" CAI>/agent 1 2 3", "green"), " - Switch to agent by position number\n", + (" CAI>/agent", "green"), " - Display list of all available agents\n\n", + ("MODEL SELECTION", "bold yellow"), "\n", + (" CAI>/model [NAME]", "green"), " - Change to a different model by name\n", + (" CAI>/model 1", "green"), " - Change model by position number\n", + (" CAI>/model", "green"), " - Show all available models\n\n", + ("INPUT & EXECUTION", "bold yellow"), "\n", + (" ESC + ENTER", "green"), " - Enter multi-line input mode\n", + (" CAI>/shell or CAI> $", "green"), " - Run system shell commands\n", + (" CAI>hi, cybersecurity AI", "green"), " - Any text without commands will be sent as a prompt\n", + (" CAI>/help", "green"), " - Display complete command reference\n", + (" CAI>/flush or CAI> /clear", "green"), " - Clear the conversation history\n\n", + ("UTILITY COMMANDS", "bold yellow"), "\n", + (" CAI>/mcp", "green"), " - Load additional tools with MCP server to an agent\n", + (" CAI>/virt", "green"), " - Show all available virtualized environments\n", + (" CAI>/flush", "green"), " - Flush context/message list\n", + ("━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━", "dim"), "\n", + ) + + # Get current environment variable values + current_model = os.getenv('CAI_MODEL', "qwen2.5:14b") + current_agent_type = os.getenv('CAI_AGENT_TYPE', "one_tool_agent") + + config_text = Text.assemble( + ("Quick Start Configuration", "bold cyan underline"), "\n\n", + ("1. Configure .env file with your settings", "yellow"), "\n", + ("2. Select an agent: ", "yellow"), f"by default: CAI_AGENT_TYPE={current_agent_type}\n", + ("3. Select a model: ", "yellow"), f"by default: CAI_MODEL={current_model}\n\n", + ("Basic Usage:", "bold yellow"), "\n", + (" 1. CAI> /model", "green"), " - View all available models first\n", + (" 2. CAI> /agent", "green"), " - View all available agents first\n", + (" 3. CAI> /model deepseek/deepseek-chat", "green"), " - Then select your preferred model\n", + (" 4. CAI> /agent 16", "green"), " - Then select your preferred agent\n", + (" 5. CAI> Scan 192.168.1.1", "green"), " - Example prompt for target scan\n\n", + (" /help", "green"), " - Display complete command reference\n\n", + ("Common Environment Variables:", "bold yellow"), "\n", + (" CAI_MODEL", "green"), f" - Model to use (default: {current_model})\n", + (" CAI_AGENT_TYPE", "green"), f" - Agent type (default: {current_agent_type})\n", + (" CAI_DEBUG", "green"), f" - Debug level (default: {os.getenv('CAI_DEBUG', '1')})\n", + (" CAI_MAX_TURNS", "green"), f" - Max conversation turns (default: {os.getenv('CAI_MAX_TURNS', 'inf')})\n", + (" CAI_TRACING", "green"), f" - Enable tracing (default: {os.getenv('CAI_TRACING', 'true')})\n", + ) + + # Create additional tips panels + ollama_tip = Panel( + "To use Ollama models, configure OLLAMA_API_BASE\n" + "before startup.\n\n" + "Default: host.docker.internal:8000/v1", + title="[bold yellow]Ollama Configuration[/bold yellow]", + border_style="yellow", + padding=(1, 2), + title_align="center" + ) + + context_tip = Panel( + "As security exercises progress, LLM quality may\n" + "degrade, especially if progress stalls.\n\n" + "It's often better to clear the context window\n" + "or restart CAI rather than waiting until\n" + "context usage reaches 100%.\n\n" + "When context exceeds 80%, follow these steps:\n" + "1. CAI> Dump your memory and findings in current scenario in findings.txt\n" + "2. CAI> /flush\n" + "3. CAI> Analyze findings.txt, and continue exercise with target: ...", + title="[bold yellow]Performance Tip[/bold yellow]", + border_style="yellow", + padding=(1, 2), + title_align="center" + ) + + # Combine tips into a group + # tips_group = Group(ollama_tip, context_tip) + tips_group = Group(context_tip) + + # Create a three-column panel layout + console.print(Panel( + Columns( + [help_text, config_text, tips_group], + column_first=True, + expand=True, + align="center" + ), + title="[bold]CAI Quick Guide[/bold]", + border_style="blue", + padding=(1, 2), + title_align="center" + )) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index e294bee7..08d5dfe2 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -1632,17 +1632,31 @@ class _Converter: model_instance = self_obj break - token_info = {} - if model_instance: - token_info = { - 'interaction_input_tokens': getattr(model_instance, 'interaction_input_tokens', 0), - 'interaction_output_tokens': getattr(model_instance, 'interaction_output_tokens', 0), - 'interaction_reasoning_tokens': getattr(model_instance, 'interaction_reasoning_tokens', 0), - 'total_input_tokens': getattr(model_instance, 'total_input_tokens', 0), - 'total_output_tokens': getattr(model_instance, 'total_output_tokens', 0), - 'total_reasoning_tokens': getattr(model_instance, 'total_reasoning_tokens', 0), - 'model': str(getattr(model_instance, 'model', '')), - } + # Always create a token_info dictionary, even if some values are zero + token_info = { + 'interaction_input_tokens': getattr(model_instance, 'interaction_input_tokens', 0), + 'interaction_output_tokens': getattr(model_instance, 'interaction_output_tokens', 0), + 'interaction_reasoning_tokens': getattr(model_instance, 'interaction_reasoning_tokens', 0), + 'total_input_tokens': getattr(model_instance, 'total_input_tokens', 0), + 'total_output_tokens': getattr(model_instance, 'total_output_tokens', 0), + 'total_reasoning_tokens': getattr(model_instance, 'total_reasoning_tokens', 0), + 'model': str(getattr(model_instance, 'model', '')), + } + + # Calculate costs using standard cost model + if model_instance and hasattr(model_instance, 'model'): + from cai.util import calculate_model_cost + model_name = str(model_instance.model) + token_info['interaction_cost'] = calculate_model_cost( + model_name, + token_info['interaction_input_tokens'], + token_info['interaction_output_tokens'] + ) + token_info['total_cost'] = calculate_model_cost( + model_name, + token_info['total_input_tokens'], + token_info['total_output_tokens'] + ) # Use the cli_print_tool_output function with actual token values cli_print_tool_output( diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 91d9dbbe..e1e58778 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -12,6 +12,14 @@ import time import uuid import sys from wasabi import color # pylint: disable=import-error +from cai.util import format_time + + +# Instead of direct import +try: + from cai.cli import START_TIME +except ImportError: + START_TIME = None # Global dictionary to store active sessions ACTIVE_SESSIONS = {} @@ -228,10 +236,10 @@ def _run_ctf(ctf, command, stdout=False, timeout=100, stream=False, call_id=None return f"Error executing CTF command: {str(e)}" -def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None): +def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None): # If streaming is enabled and we have a call_id if stream and call_id: - return _run_local_streamed(command, call_id, timeout) + return _run_local_streamed(command, call_id, timeout, tool_name) try: # nosec B602 - shell=True is required for command chaining @@ -265,7 +273,7 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None): return error_msg -def _run_local_streamed(command, call_id, timeout=100): +def _run_local_streamed(command, call_id, timeout=100, tool_name=None): """Run a local command with streaming output to the Tool output panel""" try: # Try to import Rich for nice display @@ -293,23 +301,43 @@ def _run_local_streamed(command, call_id, timeout=100): bufsize=1 ) + # If tool_name is not provided, derive it from the command + if tool_name is None: + # Just use the first command as the tool name + tool_name = command.strip().split()[0] + "_command" + # Create panel content for Rich display if rich_available: - tool_name = "generic_linux_command" # Parse command into command and args parts = command.strip().split(' ', 1) cmd = parts[0] if parts else "" args = parts[1] if len(parts) > 1 else "" + # Format clean arguments, following the same rules as cli_print_tool_output + arg_parts = [] + if cmd: + arg_parts.append(f"command={cmd}") + if args and args.strip(): # Only add args if non-empty + arg_parts.append(f"args={args}") + args_str = ", ".join(arg_parts) + header = Text() header.append(tool_name, style="#00BCD4") header.append("(", style="yellow") - # Format to match: generic_linux_command({"command":"ls","args":"-la","ctf":{},"async_mode":false,"session_id":""}) - header.append(f'{{"command":"{cmd}","args":"{args}","ctf":{{}},"async_mode":false,"session_id":""}}', style="yellow") + header.append(args_str, style="yellow") header.append(")", style="yellow") - + tool_time = 0 + start_time = time.time() + total_time = time.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() - content.append(f"Executing: {command}\n\n", style="green") panel = Panel( Text.assemble(header, "\n\n", content), @@ -323,15 +351,35 @@ def _run_local_streamed(command, call_id, timeout=100): # 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 + # 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]", @@ -341,7 +389,6 @@ def _run_local_streamed(command, call_id, timeout=100): box=ROUNDED ) live.update(panel) - # Check if process is done process.stdout.close() return_code = process.wait(timeout=timeout) @@ -363,12 +410,9 @@ def _run_local_streamed(command, call_id, timeout=100): live.update(panel) # Add completion message - completion_status = "Completed" if return_code == 0 else f"Failed (code {return_code})" - content.append(f"\nCommand {completion_status}", style="green") panel = Panel( Text.assemble(header, "\n\n", content), title="[bold green]Tool Execution[/bold green]", - subtitle=f"[bold green]{completion_status}[/bold green]", border_style="green", padding=(1, 2), box=ROUNDED @@ -383,10 +427,17 @@ def _run_local_streamed(command, call_id, timeout=100): parts = command.strip().split(' ', 1) cmd = parts[0] if parts else "" args = parts[1] if len(parts) > 1 else "" - tool_args = {"command": cmd, "args": args, "ctf": {}, "async_mode": False, "session_id": ""} + + # Create a dictionary with only non-empty values (following the same rules) + tool_args = {} + if cmd: + tool_args["command"] = cmd + if args and args.strip(): + tool_args["args"] = args + # Note: Omitted empty values and async_mode=False as it's default # Initial notification - just once - cli_print_tool_output("generic_linux_command", tool_args, "Command started...", call_id=call_id) + cli_print_tool_output(tool_name, tool_args, "Command started...", call_id=call_id) # Buffer for collecting output buffer_size = 0 @@ -404,7 +455,7 @@ def _run_local_streamed(command, call_id, timeout=100): # Only update the output periodically to reduce panel refresh rate if buffer_size >= update_interval: current_output = ''.join(output_buffer) - cli_print_tool_output("generic_linux_command", tool_args, current_output, call_id=call_id) + cli_print_tool_output(tool_name, tool_args, current_output, call_id=call_id) buffer_size = 0 # Check if process is done @@ -421,7 +472,7 @@ def _run_local_streamed(command, call_id, timeout=100): if return_code != 0: final_output += f"\nCommand exited with code {return_code}" - cli_print_tool_output("generic_linux_command", tool_args, final_output, call_id=call_id) + cli_print_tool_output(tool_name, tool_args, final_output, call_id=call_id) # Return the full output return ''.join(output_buffer) @@ -434,7 +485,7 @@ def _run_local_streamed(command, call_id, timeout=100): # Update tool output panel with timeout message if not rich_available: tool_args = {"command": command} - cli_print_tool_output("generic_linux_command", tool_args, final_output, call_id=call_id) + cli_print_tool_output(tool_name, tool_args, final_output, call_id=call_id) return final_output @@ -445,14 +496,14 @@ def _run_local_streamed(command, call_id, timeout=100): # Update tool output panel with error message if simple streaming if not rich_available: tool_args = {"command": command} - cli_print_tool_output("generic_linux_command", tool_args, error_msg, call_id=call_id) + cli_print_tool_output(tool_name, tool_args, error_msg, call_id=call_id) return error_msg def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arguments # noqa: E501 async_mode=False, session_id=None, - timeout=100, stream=False, call_id=None): + timeout=100, stream=False, call_id=None, tool_name=None): """ Run command either in CTF container or on the local attacker machine @@ -465,6 +516,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg timeout: Command timeout in seconds stream: Whether to stream output in real-time call_id: Unique ID for the command execution (for streaming) + tool_name: Name of the tool being executed (for display in streaming output). + If None, the tool name will be derived from the command. Returns: str: Command output, status message, or session ID @@ -497,4 +550,4 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg # Otherwise, run command normally if ctf: return _run_ctf(ctf, command, stdout, timeout, stream, call_id) - return _run_local(command, stdout, timeout, stream, call_id) + return _run_local(command, stdout, timeout, stream, call_id, tool_name) diff --git a/src/cai/tools/reconnaissance/generic_linux_command.py b/src/cai/tools/reconnaissance/generic_linux_command.py index 6c1274b8..81fbf0ba 100644 --- a/src/cai/tools/reconnaissance/generic_linux_command.py +++ b/src/cai/tools/reconnaissance/generic_linux_command.py @@ -99,10 +99,11 @@ def generic_linux_command(command: str = "", if stream and not call_id: call_id = str(uuid.uuid4())[:8] - # Run the command with the appropriate parameters + # Run the command with the appropriate parameters - pass the actual tool name! result = run_command(full_command, ctf=ctf, async_mode=async_mode, session_id=session_id, - timeout=timeout, stream=stream, call_id=call_id) + timeout=timeout, stream=stream, call_id=call_id, + tool_name="generic_linux_command") # For better output formatting, if we're in streaming mode, we can modify the output # to include any additional information needed while still preventing the duplicate panel diff --git a/src/cai/util.py b/src/cai/util.py index 087c98d2..63cedaf5 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -19,7 +19,13 @@ from datetime import datetime import atexit from dataclasses import dataclass, field from typing import Dict, Optional +import time +# Instead of direct import +try: + from cai.cli import START_TIME +except ImportError: + START_TIME = None # Shared stats tracking object to maintain consistent costs across calls @dataclass @@ -521,7 +527,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. @@ -831,6 +852,13 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli text.append(f" ({os.getenv('CAI_SUPPORT_MODEL')})", style="bold blue") text.append("]", style="dim") + elif not parsed_message: + # When parsed_message is empty, only include timestamp and model info + text.append(f"Agent: {agent_name} ", style="bold green") + text.append(f"[{timestamp}", style="dim") + if model: + text.append(f" ({model})", style="bold magenta") + text.append("]", style="dim") else: text.append(f"[{counter}] ", style="bold cyan") text.append(f"Agent: {agent_name} ", style="bold green") @@ -861,7 +889,9 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli interaction_cost, total_cost ) - text.append(tokens_text) + # Only append token information if there is a parsed message + if parsed_message: + text.append(tokens_text) panel = Panel( text, @@ -873,7 +903,7 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli else "[bold]Agent Interaction[/bold]"), title_align="left" ) - console.print("\n") + #console.print("\n") console.print(panel) # If there are tool panels, print them after the main message panel @@ -1052,7 +1082,6 @@ def finish_agent_streaming(context, final_stats=None): context["live"].update(final_panel) # Ensure updates are displayed before stopping - import time time.sleep(0.5) # Stop the live display @@ -1069,12 +1098,17 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut output: The output of the tool call_id: Optional call ID for streaming updates execution_info: Optional execution information - token_info: Optional token information + token_info: Optional token information with keys: + - interaction_input_tokens, interaction_output_tokens, interaction_reasoning_tokens + - total_input_tokens, total_output_tokens, total_reasoning_tokens + - model: model name string + - interaction_cost, total_cost: optional cost values """ # If it's an empty output, don't print anything if not output and not call_id: return + # CRITICAL CHECK: When in streaming mode (CAI_STREAM=true), ONLY show output panels # for streaming updates (those with call_id). This prevents duplicate output panels. is_streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true' @@ -1112,33 +1146,130 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut from rich.panel import Panel from rich.text import Text from rich.box import ROUNDED + from rich.console import Group # Create a console for output - console = Console() + console = Console(theme=theme) - # Format arguments as a string if they are a dictionary + # Format arguments for display + # Parse JSON string if args is a string + if isinstance(args, str) and args.strip().startswith('{'): + try: + import json + args = json.loads(args) + except: + # Keep as is if not valid JSON + pass + + # Format arguments as a clean string if isinstance(args, dict): - args_str = ", ".join([f"{key}='{value}'" for key, value in args.items()]) + # Only include non-empty values and exclude async_mode=false + arg_parts = [] + for key, value in args.items(): + # Skip empty values + if value == "" or value == {} or value is None: + continue + # Skip async_mode=false (default) + if key == "async_mode" and value is False: + continue + # Format the value + if isinstance(value, str): + arg_parts.append(f"{key}={value}") + else: + arg_parts.append(f"{key}={value}") + args_str = ", ".join(arg_parts) else: args_str = str(args) + # Get session timing information + try: + from cai.cli import START_TIME + total_time = time.time() - START_TIME if START_TIME else None + except ImportError: + total_time = None + + + # Extract execution timing info + + tool_time = None + status = None + if execution_info: + # Prefer 'tool_time' if present, else fallback to 'time_taken' + + tool_time = execution_info.get('tool_time') + status = execution_info.get('status', 'completed') + + # Create header for all panel displays (both streaming and non-streaming) + header = Text() + header.append(tool_name, style="#00BCD4") + header.append("(", style="yellow") + header.append(args_str, style="yellow") + header.append(")", style="yellow") + + # Add timing information directly in the 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)}") + if timing_info: + header.append(f" [{' | '.join(timing_info)}]", style="cyan") + + # Add completion status if available - REMOVED, just showing timing now + # if status: + # if status == 'completed': + # header.append(f" [Completed]", style="green") + # elif status == 'running': + # header.append(f" [Running]", style="yellow") + # elif status == 'error': + # header.append(f" [Error]", style="red") + # elif status == 'timeout': + # header.append(f" [Timeout]", style="red") + # else: + # header.append(f" [{status.title()}]", style="dim") + # For streaming mode with call_id, use Rich Live display if call_id: - # Create the header text - header = Text() - header.append(tool_name, style="#00BCD4") - header.append("(", style="yellow") - header.append(args_str, style="yellow") - header.append(")", style="yellow") + # Create token information if available + token_content = None + if token_info: + model = token_info.get('model', '') + interaction_input_tokens = token_info.get('interaction_input_tokens', 0) + interaction_output_tokens = token_info.get('interaction_output_tokens', 0) + interaction_reasoning_tokens = token_info.get('interaction_reasoning_tokens', 0) + total_input_tokens = token_info.get('total_input_tokens', 0) + total_output_tokens = token_info.get('total_output_tokens', 0) + total_reasoning_tokens = token_info.get('total_reasoning_tokens', 0) + + if (interaction_input_tokens > 0 or total_input_tokens > 0): + token_text = _create_token_display( + interaction_input_tokens, + interaction_output_tokens, + interaction_reasoning_tokens, + total_input_tokens, + total_output_tokens, + total_reasoning_tokens, + model, + token_info.get('interaction_cost'), + token_info.get('total_cost') + ) + token_content = Text("\n\n") + token_content.append(token_text) # Create content text with the output content = Text(output) - # Create the panel for display + # Create the panel for display, including token info if available + panel_content = [header, Text("\n\n"), content] + if token_content: + panel_content.append(token_content) + + # Create title - simple title with no timing info + title = "[bold blue]Tool Output[/bold blue]" + panel = Panel( - Text.assemble(header, "\n\n", content), - title=f"[bold blue]Tool Output[/bold blue]", #(ID: {call_id})[/bold green]", - #subtitle="[bold green]Live Update[/bold green]", + Text.assemble(*panel_content), + title=title, border_style="blue", padding=(1, 2), box=ROUNDED, @@ -1149,52 +1280,59 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut console.print(panel) return - # For non-streaming output, show a standard panel - tool_call = f"{tool_name}({args_str})" - - # Add execution info if available - subtitle = "[bold green]Completed[/bold green]" - if execution_info: - time_taken = execution_info.get('time_taken', 0) - status = execution_info.get('status', 'completed') - - if time_taken: - subtitle = f"[bold green]{status.title()} in {time_taken:.2f}s[/bold green]" - else: - subtitle = f"[bold green]{status.title()}[/bold green]" - - # Create content sections - sections = [] - - # Add token info if available + # For non-streaming output, also use a blue panel with the same format + # Create token information if available + token_text = None if token_info: - tokens_in = token_info.get('input_tokens', 0) - tokens_out = token_info.get('output_tokens', 0) - cost = token_info.get('cost', 0) + model = token_info.get('model', '') + interaction_input_tokens = token_info.get('interaction_input_tokens', 0) + interaction_output_tokens = token_info.get('interaction_output_tokens', 0) + interaction_reasoning_tokens = token_info.get('interaction_reasoning_tokens', 0) + total_input_tokens = token_info.get('total_input_tokens', 0) + total_output_tokens = token_info.get('total_output_tokens', 0) + total_reasoning_tokens = token_info.get('total_reasoning_tokens', 0) + interaction_cost = token_info.get('interaction_cost') + total_cost = token_info.get('total_cost') - token_text = Text() - if tokens_in or tokens_out: - token_text.append(f"Tokens: {tokens_in} in, {tokens_out} out", style="cyan") - if cost: - token_text.append(f"\nCost: ${cost:.6f}", style="cyan") - - if token_text: - sections.append(token_text) + # Generate token display with CostTracker + if (interaction_input_tokens > 0 or total_input_tokens > 0): + token_text = _create_token_display( + interaction_input_tokens, + interaction_output_tokens, + interaction_reasoning_tokens, + total_input_tokens, + total_output_tokens, + total_reasoning_tokens, + model, + interaction_cost, + total_cost + ) - # Add the main output + # Now create the panel content, starting with the header + panel_content = [header, Text("\n")] + + # Add token display if available + if token_text: + panel_content.append(token_text) + panel_content.append(Text("\n")) # Add spacing after token display + + # Add the output if output: - output_text = Text() - if sections: # Add a separator if we have previous sections - output_text.append("\n") - output_text.append(output) - sections.append(output_text) + output_text = Text(output) + panel_content.append(output_text) - # Create the panel with all sections + # If no content was added but we have output, add it directly + if len(panel_content) == 2 and output: # Only header and newline + panel_content.append(Text(output)) + + # Create title - simple title with no timing info + title = "[bold blue]Tool Output[/bold blue]" + + # Create the final panel - always blue now panel = Panel( - Text.assemble(*sections), - title=f"[bold green]Tool Output: {tool_call}[/bold green]", - subtitle=subtitle, - border_style="green", + Group(*panel_content), + title=title, + border_style="blue", padding=(1, 2), box=ROUNDED ) @@ -1204,64 +1342,112 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut except ImportError: # Fall back to simple formatting if Rich is not available - # Format arguments as a string if they are a dictionary + # Format arguments in the cleaner format + # Parse JSON string if args is a string + if isinstance(args, str) and args.strip().startswith('{'): + try: + import json + args = json.loads(args) + except: + # Keep as is if not valid JSON + pass + + # Format arguments as a clean string if isinstance(args, dict): - args_str = ", ".join([f"{key}='{value}'" for key, value in args.items()]) + # Only include non-empty values and exclude async_mode=false + arg_parts = [] + for key, value in args.items(): + # Skip empty values + if value == "" or value == {} or value is None: + continue + # Skip async_mode=false (default) + if key == "async_mode" and value is False: + continue + # Format the value + if isinstance(value, str): + arg_parts.append(f"{key}={value}") + else: + arg_parts.append(f"{key}={value}") + args_str = ", ".join(arg_parts) else: args_str = str(args) - - # Simplify output presentation for streaming mode - if call_id: - # This is a streaming update, so we need to overwrite previous output - # We'll use a basic format that's better suited for streaming + # Get session timing information + try: + from cai.cli import START_TIME + total_time = time.time() - START_TIME if START_TIME else None + except ImportError: + total_time = None - # Get terminal width for better formatting - try: - term_width = os.get_terminal_size().columns - except: # pylint: disable=bare-except - term_width = 80 - - # Create a header for the tool output - header = f"{tool_name}({args_str})" - header = header[:term_width-4] - - # Clear the screen for the tool output (alternative approach) - print(f"\r{header}") - - # Print the content - # For streaming updates, we'll use a simple format - print(output) - - # Force flush to ensure output is displayed immediately - sys.stdout.flush() - return # For non-streaming output, use the original formatting tool_call = f"{tool_name}({args_str})" - # If there's execution info, add it to the output + # Get tool execution time if available + tool_time_str = "" + execution_status = "" if execution_info: time_taken = execution_info.get('time_taken', 0) status = execution_info.get('status', 'completed') # Add execution info to the tool call display if time_taken: - tool_call += f" [{status} in {time_taken:.2f}s]" + tool_time_str = f"Tool: {format_time(time_taken)}" + execution_status = f" [{status} in {time_taken:.2f}s]" else: - tool_call += f" [{status}]" + execution_status = f" [{status}]" - print(color(f"Tool Output: {tool_call}", fg="blue")) - - # If we have token info, display it - if token_info: - tokens_in = token_info.get('input_tokens', 0) - tokens_out = token_info.get('output_tokens', 0) - cost = token_info.get('cost', 0) + # Create timing display string + 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") - if tokens_in or tokens_out: - print(color(f" Tokens: {tokens_in} in, {tokens_out} out", fg="cyan")) - if cost: - print(color(f" Cost: ${cost:.6f}", fg="cyan")) + timing_display = f" [{' | '.join(timing_info)}]" if timing_info else "" + + # Show tool name, args, execution status and timing display + print(color(f"Tool Output: {tool_call}{timing_display}{execution_status}", fg="blue")) + + # If we have token info, display it using the consistent format from _create_token_display + if token_info: + model = token_info.get('model', '') + interaction_input_tokens = token_info.get('interaction_input_tokens', 0) + interaction_output_tokens = token_info.get('interaction_output_tokens', 0) + interaction_reasoning_tokens = token_info.get('interaction_reasoning_tokens', 0) + total_input_tokens = token_info.get('total_input_tokens', 0) + total_output_tokens = token_info.get('total_output_tokens', 0) + total_reasoning_tokens = token_info.get('total_reasoning_tokens', 0) + interaction_cost = token_info.get('interaction_cost') + total_cost = token_info.get('total_cost') + + # If we have complete token information, display it + if (interaction_input_tokens > 0 or total_input_tokens > 0): + # Manually create formatted output similar to _create_token_display + print(color(f" Current: I:{interaction_input_tokens} O:{interaction_output_tokens} R:{interaction_reasoning_tokens}", fg="cyan")) + + # Calculate or use provided costs + current_cost = COST_TRACKER.process_interaction_cost( + model, + interaction_input_tokens, + interaction_output_tokens, + interaction_reasoning_tokens, + interaction_cost + ) + total_cost_value = COST_TRACKER.process_total_cost( + model, + total_input_tokens, + total_output_tokens, + total_reasoning_tokens, + total_cost + ) + print(color(f" Cost: Current ${current_cost:.4f} | Total ${total_cost_value:.4f} | Session ${COST_TRACKER.session_total_cost:.4f}", fg="cyan")) + + # Show context usage + context_pct = interaction_input_tokens / get_model_input_tokens(model) * 100 + indicator = "🟩" if context_pct < 50 else "🟨" if context_pct < 80 else "🟥" + print(color(f" Context: {context_pct:.1f}% {indicator}", fg="cyan")) # Print the actual output print(output) diff --git a/tests/tools/test_tool_generic_linux_command.py b/tests/tools/test_tool_generic_linux_command.py index c4c427c2..e4852f8e 100644 --- a/tests/tools/test_tool_generic_linux_command.py +++ b/tests/tools/test_tool_generic_linux_command.py @@ -1,11 +1,11 @@ -""" -This module contains tests for the generic Linux command tool functionality. -It includes tests for executing regular commands and handling invalid commands. -""" - +import os import pytest import json from unittest.mock import MagicMock + +# Set test environment variables to avoid OpenAI client initialization errors +os.environ["OPENAI_API_KEY"] = "test_key_for_ci_environment" + from cai.tools.reconnaissance.generic_linux_command import generic_linux_command async def test_generic_linux_command_regular_commands(): @@ -48,4 +48,4 @@ async def test_generic_linux_command_invalid_command(): result = await generic_linux_command.on_invoke_tool(mock_ctx, json.dumps(params)) # Assert that the result indicates the command was not found - assert "not found" in result + assert "not found" in result \ No newline at end of file