diff --git a/src/cai/cli.py b/src/cai/cli.py index 677914db..72940444 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -184,8 +184,8 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= # Function to get the short name of the agent for display def get_agent_short_name(agent): if hasattr(agent, 'name'): - # Split by spaces and take first word to get shortened name - return agent.name.split()[0] + # Return the full agent name instead of just the first word + return agent.name return "Agent" # Prevent the model from using its own rich streaming to avoid conflicts @@ -209,7 +209,9 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= get_toolbar_with_refresh, current_text ) - + except KeyboardInterrupt: + break + try: # Handle special commands if user_input.startswith('/') or user_input.startswith('$'): parts = user_input.strip().split() @@ -355,17 +357,17 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= # Use non-streamed response console.print("[dim]Thinking...[/dim]") response = asyncio.run(Runner.run(agent, user_input)) - console.print(f"Agent: {response.final_output}") + #console.print(f"Agent: {response.final_output}") # NOTE: this line is commented to avoid duplicate output turn_count += 1 except KeyboardInterrupt: - # Ensure streaming context is cleaned up on keyboard interrupt - if current_streaming_context is not None: - try: - current_streaming_context["live"].stop() - except Exception: - pass - current_streaming_context = None - break + if stream: + # Ensure streaming context is cleaned up on keyboard interrupt + if current_streaming_context is not None: + try: + current_streaming_context["live"].stop() + except Exception: + pass + current_streaming_context = None except Exception as e: # Ensure streaming context is cleaned up on any exception if current_streaming_context is not None: diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 4031a7af..e294bee7 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -6,6 +6,7 @@ import time import os import litellm import tiktoken +import inspect from collections.abc import AsyncIterator, Iterable from dataclasses import dataclass, field @@ -288,28 +289,51 @@ class OpenAIChatCompletionsModel(Model): hasattr(response.usage.completion_tokens_details, 'reasoning_tokens')): self.total_reasoning_tokens += response.usage.completion_tokens_details.reasoning_tokens - # Print the agent message for CLI display - cli_print_agent_messages( - agent_name=getattr(self, 'agent_name', 'Agent'), # Default to 'Agent' if not available - message=response.choices[0].message, - counter=getattr(self, 'interaction_counter', 0), # Default to 0 if not available - model=str(self.model), - debug=False, - interaction_input_tokens=input_tokens, - interaction_output_tokens=output_tokens, - interaction_reasoning_tokens=( - response.usage.completion_tokens_details.reasoning_tokens - if response.usage and hasattr(response.usage, 'completion_tokens_details') - and response.usage.completion_tokens_details - and hasattr(response.usage.completion_tokens_details, 'reasoning_tokens') - else 0 - ), - total_input_tokens=getattr(self, 'total_input_tokens', 0), # Will need to be tracked elsewhere - total_output_tokens=getattr(self, 'total_output_tokens', 0), # Will need to be tracked elsewhere - total_reasoning_tokens=getattr(self, 'total_reasoning_tokens', 0), # Will need to be tracked elsewhere - interaction_cost=None, # Would need cost calculation logic - total_cost=None, # Would need cost calculation logic - ) + # Check if this message contains tool calls + tool_output = None + should_display_message = True + + if (hasattr(response.choices[0].message, 'tool_calls') and + response.choices[0].message.tool_calls): + + # For each tool call in the message, get corresponding output if available + for tool_call in response.choices[0].message.tool_calls: + call_id = tool_call.id + + # If we're using direct tool output display with cli_print_tool_output, + # and we've already displayed this tool call output, we can skip displaying + # the assistant message to avoid duplication + if (hasattr(_Converter, 'tool_outputs') and call_id in _Converter.tool_outputs and + hasattr(_Converter, 'recent_tool_calls') and call_id in _Converter.recent_tool_calls): + # We've already displayed this tool and its output directly + should_display_message = False + break + + # Only display the agent message if we haven't already shown the tool output + if should_display_message: + # Print the agent message for CLI display + cli_print_agent_messages( + agent_name=getattr(self, 'agent_name', 'Agent'), + message=response.choices[0].message, + counter=getattr(self, 'interaction_counter', 0), + model=str(self.model), + debug=False, + interaction_input_tokens=input_tokens, + interaction_output_tokens=output_tokens, + interaction_reasoning_tokens=( + response.usage.completion_tokens_details.reasoning_tokens + if response.usage and hasattr(response.usage, 'completion_tokens_details') + and response.usage.completion_tokens_details + and hasattr(response.usage.completion_tokens_details, 'reasoning_tokens') + else 0 + ), + total_input_tokens=getattr(self, 'total_input_tokens', 0), + total_output_tokens=getattr(self, 'total_output_tokens', 0), + total_reasoning_tokens=getattr(self, 'total_reasoning_tokens', 0), + interaction_cost=None, + total_cost=None, + tool_output=None, # Don't pass tool output here, we're using direct display + ) usage = ( Usage( @@ -823,6 +847,22 @@ class OpenAIChatCompletionsModel(Model): "output_tokens": output_tokens, } + # To avoid duplicate tool output display, we need to track tool calls + # Add this after the completion response is received + + if not stream and hasattr(response, 'choices') and len(response.choices) > 0: + # For non-streaming responses, make sure we capture tool call IDs + # to prevent duplicate printing + choice = response.choices[0] + if hasattr(choice, 'message') and hasattr(choice.message, 'tool_calls'): + for tool_call in choice.message.tool_calls: + if hasattr(tool_call, 'id'): + # Register this tool call ID as already seen + from cai.util import cli_print_tool_output + if not hasattr(cli_print_tool_output, '_seen_calls'): + cli_print_tool_output._seen_calls = {} + cli_print_tool_output._seen_calls[tool_call.id] = True + @overload async def _fetch_response( self, @@ -1507,6 +1547,23 @@ class _Converter: elif func_call := cls.maybe_function_tool_call(item): asst = ensure_assistant_message() tool_calls = list(asst.get("tool_calls", [])) + + # Save the tool call details for later matching with output + if not hasattr(cls, 'recent_tool_calls'): + cls.recent_tool_calls = {} + + # Store the tool call by ID for later reference + # Also store the current time for execution timing + import time + cls.recent_tool_calls[func_call["call_id"]] = { + 'name': func_call["name"], + 'arguments': func_call["arguments"], + 'start_time': time.time(), + 'execution_info': { + 'start_time': time.time() + } + } + new_tool_call = ChatCompletionMessageToolCallParam( id=func_call["call_id"], type="function", @@ -1517,8 +1574,87 @@ class _Converter: ) tool_calls.append(new_tool_call) asst["tool_calls"] = tool_calls + # 5) function call output => tool message elif func_output := cls.maybe_function_tool_call_output(item): + # Store the output for this call_id + call_id = func_output["call_id"] + output_content = func_output["output"] + + # Update execution timing if we have the start time + if hasattr(cls, 'recent_tool_calls') and call_id in cls.recent_tool_calls: + tool_call = cls.recent_tool_calls[call_id] + if 'start_time' in tool_call: + end_time = time.time() + tool_execution_time = end_time - tool_call['start_time'] + + # Update the execution info + if 'execution_info' in tool_call: + tool_call['execution_info']['end_time'] = end_time + tool_call['execution_info']['tool_time'] = tool_execution_time + + # If this is the first tool being executed, record the total time from conversation start + if not hasattr(cls, 'conversation_start_time'): + cls.conversation_start_time = tool_call['start_time'] + + total_time = end_time - getattr(cls, 'conversation_start_time', tool_call['start_time']) + tool_call['execution_info']['total_time'] = total_time + + # Store the output so it can be accessed later + if not hasattr(cls, 'tool_outputs'): + cls.tool_outputs = {} + + cls.tool_outputs[call_id] = output_content + + # Display the tool output immediately with the matched tool call + from cai.util import cli_print_tool_output + + # Check if we're in streaming mode - don't show tool output panel in streaming mode + is_streaming_enabled = os.environ.get('CAI_STREAM', 'false').lower() == 'true' + if is_streaming_enabled: + # Don't display tool output in streaming mode - it will be handled elsewhere + pass # Just skip the display, but preserve the tool output + else: + # For non-streaming mode, maintain the original behavior + # Look up the original tool call to get the name and arguments + if hasattr(cls, 'recent_tool_calls') and call_id in cls.recent_tool_calls: + tool_call = cls.recent_tool_calls[call_id] + tool_name = tool_call.get('name', 'Unknown Tool') + tool_args = tool_call.get('arguments', {}) + execution_info = tool_call.get('execution_info', {}) + + # Get token counts from the OpenAIChatCompletionsModel if available + model_instance = None + for frame in inspect.stack(): + if 'self' in frame.frame.f_locals: + self_obj = frame.frame.f_locals['self'] + if isinstance(self_obj, OpenAIChatCompletionsModel): + 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', '')), + } + + # Use the cli_print_tool_output function with actual token values + cli_print_tool_output( + tool_name=tool_name, + args=tool_args, + output=output_content, + call_id=call_id, # Keep call_id for non-streaming mode + execution_info=execution_info, + token_info=token_info + ) + + # Continue with normal processing flush_assistant_message() msg: ChatCompletionToolMessageParam = { "role": "tool", diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 91a9752a..91d9dbbe 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -10,6 +10,7 @@ import pty import signal import time import uuid +import sys from wasabi import color # pylint: disable=import-error # Global dictionary to store active sessions @@ -211,7 +212,7 @@ def terminate_session(session_id): return result -def _run_ctf(ctf, command, stdout=False, timeout=100): +def _run_ctf(ctf, command, stdout=False, timeout=100, stream=False, call_id=None): try: # Ensure the command is executed in a shell that supports command # chaining @@ -227,7 +228,11 @@ def _run_ctf(ctf, command, stdout=False, timeout=100): return f"Error executing CTF command: {str(e)}" -def _run_local(command, stdout=False, timeout=100): +def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None): + # If streaming is enabled and we have a call_id + if stream and call_id: + return _run_local_streamed(command, call_id, timeout) + try: # nosec B602 - shell=True is required for command chaining result = subprocess.run( @@ -240,6 +245,14 @@ def _run_local(command, stdout=False, timeout=100): output = result.stdout if result.stdout else result.stderr if stdout: print("\033[32m" + output + "\033[0m") + + # Skip passing output to cli_print_tool_output when CAI_STREAM=true + # This prevents duplicate output in streaming mode + is_streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true' + if not is_streaming_enabled: + # Optional: Add cli_print_tool_output call here if needed for non-streaming + pass + return output except subprocess.TimeoutExpired as e: error_output = e.stdout.decode() if e.stdout else str(e) @@ -252,9 +265,194 @@ def _run_local(command, stdout=False, timeout=100): return error_msg +def _run_local_streamed(command, call_id, timeout=100): + """Run a local command with streaming output to the Tool output panel""" + try: + # Try to import Rich for nice display + try: + from rich.console import Console + from rich.live import Live + from rich.panel import Panel + from rich.text import Text + from rich.box import ROUNDED + console = Console() + rich_available = True + except ImportError: + rich_available = False + from cai.util import cli_print_tool_output + + output_buffer = [] + + # Start the process + process = subprocess.Popen( + command, + shell=True, # nosec B602 + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + bufsize=1 + ) + + # 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 "" + + 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(")", style="yellow") + + content = Text() + content.append(f"Executing: {command}\n\n", style="green") + + panel = Panel( + Text.assemble(header, "\n\n", content), + title="[bold green]Tool Execution[/bold green]", + subtitle="[bold green]Live Output[/bold green]", + border_style="green", + padding=(1, 2), + box=ROUNDED + ) + + # Start Live display + with Live(panel, console=console, refresh_per_second=4) as live: + # Stream stdout in real-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") + panel = Panel( + Text.assemble(header, "\n\n", content), + title="[bold green]Tool Execution[/bold green]", + subtitle="[bold green]Live Output[/bold green]", + border_style="green", + padding=(1, 2), + box=ROUNDED + ) + live.update(panel) + + # Check if process is done + process.stdout.close() + return_code = process.wait(timeout=timeout) + + # Get any stderr output + stderr_data = process.stderr.read() + if stderr_data: + content.append("\nERROR OUTPUT:\n", style="red") + content.append(stderr_data, style="red") + output_buffer.append("\nERROR OUTPUT:\n" + stderr_data) + panel = Panel( + Text.assemble(header, "\n\n", content), + title="[bold green]Tool Execution[/bold green]", + subtitle="[bold green]Live Output[/bold green]", + border_style="green", + padding=(1, 2), + box=ROUNDED + ) + 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 + ) + live.update(panel) + + # Wait a moment for the panel to be displayed properly + time.sleep(0.5) + else: + # Fallback to simpler streaming with cli_print_tool_output + # Parse command into command and args (same as rich mode) + 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": ""} + + # Initial notification - just once + cli_print_tool_output("generic_linux_command", tool_args, "Command started...", call_id=call_id) + + # Buffer for collecting output + buffer_size = 0 + update_interval = 10 # lines + + # Stream stdout in real-time + for line in iter(process.stdout.readline, ''): + if not line: + break + + # Add to output collection + output_buffer.append(line) + buffer_size += 1 + + # 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) + buffer_size = 0 + + # Check if process is done + process.stdout.close() + return_code = process.wait(timeout=timeout) + + # Get any stderr output + stderr_data = process.stderr.read() + if stderr_data: + output_buffer.append("\nERROR OUTPUT:\n" + stderr_data) + + # Final output update - always show the final result + final_output = ''.join(output_buffer) + 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) + + # Return the full output + return ''.join(output_buffer) + + except subprocess.TimeoutExpired: + error_msg = f"Command timed out after {timeout} seconds" + output_buffer.append("\n" + error_msg) + final_output = ''.join(output_buffer) + + # 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) + + return final_output + + except Exception as e: # pylint: disable=broad-except + error_msg = f"Error executing command: {str(e)}" + print(color(error_msg, fg="red")) + + # 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) + + 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): + timeout=100, stream=False, call_id=None): """ Run command either in CTF container or on the local attacker machine @@ -264,6 +462,9 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg stdout: Whether to print output to stdout async_mode: Whether to run the command asynchronously session_id: ID of an existing session to send the command to + timeout: Command timeout in seconds + stream: Whether to stream output in real-time + call_id: Unique ID for the command execution (for streaming) Returns: str: Command output, status message, or session ID @@ -287,10 +488,13 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg time.sleep(0.5) output = get_session_output(session_id, clear=False) print("\033[32m" + output + "\033[0m") - return f"Created session { - session_id}. Use this ID to interact with the session." + return f"Created session {session_id}. Use this ID to interact with the session." + # Generate a call_id if we're streaming and one wasn't provided + if stream and not call_id: + call_id = str(uuid.uuid4())[:8] + # Otherwise, run command normally if ctf: - return _run_ctf(ctf, command, stdout, timeout) - return _run_local(command, stdout, timeout) + return _run_ctf(ctf, command, stdout, timeout, stream, call_id) + return _run_local(command, stdout, timeout, stream, call_id) diff --git a/src/cai/tools/reconnaissance/generic_linux_command.py b/src/cai/tools/reconnaissance/generic_linux_command.py index b14ba4ad..6c1274b8 100644 --- a/src/cai/tools/reconnaissance/generic_linux_command.py +++ b/src/cai/tools/reconnaissance/generic_linux_command.py @@ -1,11 +1,17 @@ """ This is used to create a generic linux command. """ +import os +import time +import uuid +import subprocess +import sys from cai.tools.common import (run_command, list_shell_sessions, get_session_output, terminate_session) # pylint: disable=import-error # noqa E501 from cai.sdk.agents import function_tool +from wasabi import color # pylint: disable=import-error @function_tool @@ -74,13 +80,30 @@ def generic_linux_command(command: str = "", async_commands = ['ssh', 'python -m http.server'] async_mode = any(cmd in full_command for cmd in async_commands) - # NOTE: review this as it's a hack to get - # around the long delays with nc connections + # For SSH sessions or async commands, use different timeout if session_id: timeout = 10 else: timeout = 100 + + # Check if streaming should be enabled + stream = os.getenv('CAI_STREAM', 'false').lower() == 'true' + + # Generate a call_id for streaming if needed + # Only use call_id for streaming mode to prevent duplicate panels + # When CAI_STREAM=true, we want the Tool Execution panel but not the Tool Output panel + call_id = str(uuid.uuid4())[:8] if stream else None + + # In streaming mode, prevent displaying both panels by forcing a call_id + # This tricks cli_print_tool_output into thinking it's being called for a streaming update + if stream and not call_id: + call_id = str(uuid.uuid4())[:8] - return run_command(full_command, ctf=ctf, + # Run the command with the appropriate parameters + result = run_command(full_command, ctf=ctf, async_mode=async_mode, session_id=session_id, - timeout=timeout) + timeout=timeout, stream=stream, call_id=call_id) + + # 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 + return result diff --git a/src/cai/util.py b/src/cai/util.py index 4bab6da4..087c98d2 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -16,32 +16,208 @@ from rich.theme import Theme # pylint: disable=import-error from rich.traceback import install # pylint: disable=import-error from rich.pretty import install as install_pretty # pylint: disable=import-error # noqa: 501 from datetime import datetime +import atexit +from dataclasses import dataclass, field +from typing import Dict, Optional + +# Shared stats tracking object to maintain consistent costs across calls +@dataclass +class CostTracker: + # Session-level stats + session_total_cost: float = 0.0 + + # Current agent stats + current_agent_total_cost: float = 0.0 + current_agent_input_tokens: int = 0 + current_agent_output_tokens: int = 0 + current_agent_reasoning_tokens: int = 0 + + # Current interaction stats + interaction_input_tokens: int = 0 + interaction_output_tokens: int = 0 + interaction_reasoning_tokens: int = 0 + interaction_cost: float = 0.0 + + # Calculation cache + model_pricing_cache: Dict[str, tuple] = field(default_factory=dict) + calculated_costs_cache: Dict[str, float] = field(default_factory=dict) + + # Track the last calculation to debug inconsistencies + last_interaction_cost: float = 0.0 + last_total_cost: float = 0.0 + + def reset_interaction_stats(self): + """Reset stats for a new interaction""" + self.interaction_input_tokens = 0 + self.interaction_output_tokens = 0 + self.interaction_reasoning_tokens = 0 + self.interaction_cost = 0.0 + + def update_session_cost(self, new_cost: float) -> None: + """Add cost to session total and log the update""" + old_total = self.session_total_cost + self.session_total_cost += new_cost + + def log_final_cost(self) -> None: + """Display final cost information at exit""" + print(f"\nTotal CAI Session Cost: ${self.session_total_cost:.6f}") + + def get_model_pricing(self, model_name: str) -> tuple: + """Get and cache pricing information for a model""" + # Use the centralized function to standardize model names + model_name = get_model_name(model_name) + + # Check cache first + if model_name in self.model_pricing_cache: + return self.model_pricing_cache[model_name] + + # Fetch from LiteLLM API + LITELLM_URL = ( + "https://raw.githubusercontent.com/BerriAI/litellm/main/" + "model_prices_and_context_window.json" + ) + + try: + import requests + response = requests.get(LITELLM_URL, timeout=2) + if response.status_code == 200: + model_pricing_data = response.json() + + # Get pricing info for the model + pricing_info = model_pricing_data.get(model_name, {}) + input_cost_per_token = pricing_info.get("input_cost_per_token", 0) + output_cost_per_token = pricing_info.get("output_cost_per_token", 0) + + # Cache the results + self.model_pricing_cache[model_name] = (input_cost_per_token, output_cost_per_token) + return input_cost_per_token, output_cost_per_token + except Exception as e: + print(f" WARNING: Error fetching model pricing: {str(e)}") + + # Default values if pricing not found + default_pricing = (0, 0) + self.model_pricing_cache[model_name] = default_pricing + return default_pricing + + def calculate_cost(self, model: str, input_tokens: int, output_tokens: int, + label: Optional[str] = None, force_calculation: bool = False) -> float: + """Calculate and cache cost for a given model and token counts""" + # Standardize model name using the central function + model_name = get_model_name(model) + + # Generate a cache key + cache_key = f"{model_name}_{input_tokens}_{output_tokens}" + + # Return cached result if available (unless force_calculation is True) + if cache_key in self.calculated_costs_cache and not force_calculation: + return self.calculated_costs_cache[cache_key] + + # Get pricing information + input_cost_per_token, output_cost_per_token = self.get_model_pricing(model_name) + + # Calculate costs - use high precision for calculations + input_cost = input_tokens * input_cost_per_token + output_cost = output_tokens * output_cost_per_token + total_cost = input_cost + output_cost + + # Cache the result with full precision + self.calculated_costs_cache[cache_key] = total_cost + + return total_cost + + def process_interaction_cost(self, model: str, + input_tokens: int, + output_tokens: int, + reasoning_tokens: int = 0, + provided_cost: Optional[float] = None) -> float: + """Process and track costs for a new interaction""" + # Standardize model name + model_name = get_model_name(model) + + # Update token counts + self.interaction_input_tokens = input_tokens + self.interaction_output_tokens = output_tokens + self.interaction_reasoning_tokens = reasoning_tokens + + # Use provided cost or calculate + if provided_cost is not None and provided_cost > 0: + self.interaction_cost = float(provided_cost) + else: + self.interaction_cost = self.calculate_cost( + model_name, input_tokens, output_tokens, + label="OFFICIAL CALCULATION: Interaction") + + self.last_interaction_cost = self.interaction_cost + + return self.interaction_cost + + def process_total_cost(self, model: str, + total_input_tokens: int, + total_output_tokens: int, + total_reasoning_tokens: int = 0, + provided_cost: Optional[float] = None) -> float: + """Process and track costs for total (cumulative) usage""" + # Standardize model name + model_name = get_model_name(model) + + # Update token counts + self.current_agent_input_tokens = total_input_tokens + self.current_agent_output_tokens = total_output_tokens + self.current_agent_reasoning_tokens = total_reasoning_tokens + + # Get previous total and add current interaction cost + previous_total = self.current_agent_total_cost + + # Add the new interaction cost + if provided_cost is not None and provided_cost > 0: + # If a total cost is explicitly provided, use it + new_total_cost = float(provided_cost) + # Calculate how much was added in this interaction + cost_diff = new_total_cost - previous_total + else: + # Simply add the current interaction cost to the previous total + cost_diff = self.interaction_cost + new_total_cost = previous_total + cost_diff + + # Only add to session total if there's genuinely new cost (and it's positive) + if cost_diff > 0: + self.update_session_cost(cost_diff) + + # Update the current agent's total cost + self.current_agent_total_cost = new_total_cost + + # Track the last total for debugging + self.last_total_cost = new_total_cost + + return new_total_cost + +# Initialize the global cost tracker +COST_TRACKER = CostTracker() + +# Register exit handler for final cost display +atexit.register(COST_TRACKER.log_final_cost) theme = Theme({ - # Primary colors - Material Design inspired - "timestamp": "#00BCD4", # Cyan 500 - "agent": "#4CAF50", # Green 500 - "arrow": "#FFFFFF", # White - "content": "#ECEFF1", # Blue Grey 50 - "tool": "#F44336", # Red 500 + "timestamp": "#00BCD4", + "agent": "#4CAF50", + "arrow": "#FFFFFF", + "content": "#ECEFF1", + "tool": "#F44336", - # Secondary colors - "cost": "#009688", # Teal 500 - "args_str": "#FFC107", # Amber 500 + "cost": "#009688", + "args_str": "#FFC107", - # UI elements - "border": "#2196F3", # Blue 500 - "border_state": "#FFD700", # Yellow (Gold), complementary to Blue 500 - "model": "#673AB7", # Deep Purple 500 - "dim": "#9E9E9E", # Grey 500 - "current_token_count": "#E0E0E0", # Grey 300 - Light grey - "total_token_count": "#757575", # Grey 600 - Medium grey - "context_tokens": "#0A0A0A", # Nearly black - Very high contrast + "border": "#2196F3", + "border_state": "#FFD700", + "model": "#673AB7", + "dim": "#9E9E9E", + "current_token_count": "#E0E0E0", + "total_token_count": "#757575", + "context_tokens": "#0A0A0A", - # Status indicators - "success": "#4CAF50", # Green 500 - "warning": "#FF9800", # Orange 500 - "error": "#F44336" # Red 500 + "success": "#4CAF50", + "warning": "#FF9800", + "error": "#F44336" }) console = Console(theme=theme) @@ -331,9 +507,62 @@ def get_model_input_tokens(model): return tokens return model_tokens["gpt"] -def _create_token_display( # pylint: disable=too-many-arguments,too-many-locals,too-many-statements,too-many-branches # noqa: E501 +def get_model_name(model): + """ + Extract a string model name from various model inputs. + Centralizes model name standardization to avoid inconsistencies (e.g. avoid passing model object instead of string name). + Args: + model: String model name or model object + + Returns: + str: Standardized model name string + """ + if isinstance(model, str): + return model + # If not a string, use environment variable + return os.environ.get('CAI_MODEL', 'qwen2.5:72b') + +def get_model_pricing(model_name): + """ + Get pricing information for a model, using the CostTracker's implementation. + This is a global helper that delegates to the CostTracker instance. + + Args: + model_name: String name of the model + + Returns: + tuple: (input_cost_per_token, output_cost_per_token) + """ + # Standardize model name + model_name = get_model_name(model_name) + + # Use the CostTracker's implementation to maintain consistency and use its cache + return COST_TRACKER.get_model_pricing(model_name) + +def calculate_model_cost(model, input_tokens, output_tokens): + """ + Calculate the cost for a given model based on token usage. + + Args: + model: The model name or object + input_tokens: Number of input tokens used + output_tokens: Number of output tokens used + + Returns: + float: The calculated cost in dollars + """ + # Use the CostTracker to handle duplicates + return COST_TRACKER.calculate_cost( + model, + input_tokens, + output_tokens, + label="COST CALCULATION", + force_calculation=False # Let it use the cache for duplicates + ) + +def _create_token_display( interaction_input_tokens, - interaction_output_tokens, # noqa: E501, pylint: disable=R0913 + interaction_output_tokens, interaction_reasoning_tokens, total_input_tokens, total_output_tokens, @@ -341,60 +570,59 @@ def _create_token_display( # pylint: disable=too-many-arguments,too-many-locals model, interaction_cost=None, total_cost=None -) -> Text: # noqa: E501 - """ - Create a Text object displaying token usage information - with enhanced formatting. - """ - # print(f"\nDEBUG _create_token_display: Received costs - Interaction: {interaction_cost}, Total: {total_cost}") +) -> Text: + # Standardize model name + model_name = get_model_name(model) + # Process interaction cost + current_cost = COST_TRACKER.process_interaction_cost( + model_name, + interaction_input_tokens, + interaction_output_tokens, + interaction_reasoning_tokens, + interaction_cost + ) + + # Process total cost + total_cost_value = COST_TRACKER.process_total_cost( + model_name, + total_input_tokens, + total_output_tokens, + total_reasoning_tokens, + total_cost + ) + + # Create display text tokens_text = Text(justify="left") - - # Create a more compact, horizontal display - tokens_text.append(" ", style="bold") # Small padding + tokens_text.append(" ", style="bold") # Current interaction tokens tokens_text.append("Current: ", style="bold") tokens_text.append(f"I:{interaction_input_tokens} ", style="green") tokens_text.append(f"O:{interaction_output_tokens} ", style="red") tokens_text.append(f"R:{interaction_reasoning_tokens} ", style="yellow") - - # Current cost - only calculate if not provided - if interaction_cost is None: - interaction_cost = calculate_model_cost(model, interaction_input_tokens, interaction_output_tokens) - # Ensure interaction_cost is a float - try: - current_cost = float(interaction_cost) if interaction_cost is not None else 0.0 - except (ValueError, TypeError): - current_cost = 0.0 - tokens_text.append(f"(${current_cost:.4f}) ", style="bold") # Separator tokens_text.append("| ", style="dim") - # Total tokens + # Total tokens for this agent run tokens_text.append("Total: ", style="bold") tokens_text.append(f"I:{total_input_tokens} ", style="green") tokens_text.append(f"O:{total_output_tokens} ", style="red") tokens_text.append(f"R:{total_reasoning_tokens} ", style="yellow") - - # Total cost - only calculate if not provided - if total_cost is None: - total_cost = calculate_model_cost(model, total_input_tokens, total_output_tokens) - # Ensure total_cost is a float - try: - total_cost_value = float(total_cost) if total_cost is not None else 0.0 - except (ValueError, TypeError): - total_cost_value = 0.0 - tokens_text.append(f"(${total_cost_value:.4f}) ", style="bold") # Separator tokens_text.append("| ", style="dim") + # Session total across all agents + tokens_text.append("Session: ", style="bold magenta") + tokens_text.append(f"${COST_TRACKER.session_total_cost:.4f}", style="bold magenta") + # Context usage - context_pct = interaction_input_tokens / get_model_input_tokens(model) * 100 + tokens_text.append(" | ", style="dim") + context_pct = interaction_input_tokens / get_model_input_tokens(model_name) * 100 tokens_text.append("Context: ", style="bold") tokens_text.append(f"{context_pct:.1f}% ", style="bold") @@ -413,6 +641,143 @@ def _create_token_display( # pylint: disable=too-many-arguments,too-many-locals return tokens_text +def parse_message_content(message): + """ + Parse a message object to extract its textual content. + Only processes messages that don't have tool calls. + + Args: + message: Can be a string or a Message object with content attribute + + Returns: + str: The extracted content as a string + """ + # Check if this is a duplicate print from OpenAIChatCompletionsModel + # If message is already a string, return it + if isinstance(message, str): + return message + + # If message is a Message object with content attribute + if hasattr(message, 'content') and message.content is not None: + return message.content + + # If message is a dict with content key + if isinstance(message, dict) and 'content' in message: + return message['content'] + + # If we can't extract content, convert to string + return str(message) + +def parse_message_tool_call(message, tool_output=None): + """ + Parse a message object to extract its content and tool calls. + Displays tool calls in the format: tool_name({"command":"","args":"","ctf":{},"async_mode":false,"session_id":""}) + and shows the tool output in a separated panel. + + Args: + message: A Message object or dict with content and tool_calls attributes + tool_output: String containing the output from the tool execution + + Returns: + tuple: (content, tool_panels) where content is the message text and + tool_panels is a list of panels representing tool calls and outputs + """ + content = "" + tool_panels = [] + + # Extract the content text (LLM's inference) + if isinstance(message, str): + content = message + elif hasattr(message, 'content') and message.content is not None: + content = message.content + elif isinstance(message, dict) and 'content' in message: + content = message['content'] + + # Extract tool calls + tool_calls = None + if hasattr(message, 'tool_calls') and message.tool_calls: + tool_calls = message.tool_calls + elif isinstance(message, dict) and 'tool_calls' in message and message['tool_calls']: + tool_calls = message['tool_calls'] + + # Process tool calls if they exist + if tool_calls: + from rich.panel import Panel + from rich.text import Text + from rich.box import ROUNDED + from rich.console import Group + + for tool_call in tool_calls: + # Extract tool name and arguments + tool_name = None + args_dict = {} + call_id = None + + # Handle different formats of tool_call objects + if hasattr(tool_call, 'function'): + if hasattr(tool_call.function, 'name'): + tool_name = tool_call.function.name + if hasattr(tool_call.function, 'arguments'): + try: + import json + args_dict = json.loads(tool_call.function.arguments) + except: + args_dict = {"raw_arguments": tool_call.function.arguments} + elif isinstance(tool_call, dict): + if 'function' in tool_call: + if 'name' in tool_call['function']: + tool_name = tool_call['function']['name'] + if 'arguments' in tool_call['function']: + try: + import json + args_dict = json.loads(tool_call['function']['arguments']) + except: + args_dict = {"raw_arguments": tool_call['function']['arguments']} + + # Create a panel for this tool call if name is not None + # NOTE: Tool execution panel will be handled in cli_print_tool_output + # Pass on tool info to generate panels for display in cli_print_agent_messages + if tool_name and tool_output: + # Create content for the panel - just showing the output, not the tool call + panel_content = [] + + # Add tool output to the panel + output_text = Text() + output_text.append("Output:", style="bold #C0C0C0") # Silver/gray + output_text.append(f"\n{tool_output}", style="#C0C0C0") # Silver/gray + + panel_content.append(output_text) + + # Create a panel with just the output + tool_panel = Panel( + Group(*panel_content), + border_style="blue", + box=ROUNDED, + padding=(1, 2), + title="[bold]Tool Output[/bold]", # Changed title to indicate this is just output + title_align="left", + expand=True + ) + + tool_panels.append(tool_panel) + + # Store the call_id with tool name to help cli_print_tool_output avoid duplicates + if not hasattr(parse_message_tool_call, '_processed_calls'): + parse_message_tool_call._processed_calls = set() + + call_key = call_id if call_id else f"{tool_name}:{args_dict}" + parse_message_tool_call._processed_calls.add(call_key) + + return content, tool_panels + +# Add this function to detect tool output panels +def is_tool_output_message(message): + """Check if a message appears to be a tool output panel display message.""" + if isinstance(message, str): + msg_lower = message.lower() + return ("call id:" in msg_lower and "output:" in msg_lower) or msg_lower.startswith("tool output") + return False + def cli_print_agent_messages(agent_name, message, counter, model, debug, # pylint: disable=too-many-arguments,too-many-locals,unused-argument # noqa: E501 interaction_input_tokens=None, interaction_output_tokens=None, @@ -421,8 +786,16 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli total_output_tokens=None, total_reasoning_tokens=None, interaction_cost=None, - total_cost=None): + total_cost=None, + tool_output=None): # New parameter for tool output """Print agent messages/thoughts with enhanced visual formatting.""" + # Debug prints to trace the function calls + if debug: + if isinstance(message, str): + print(f"DEBUG cli_print_agent_messages: Received string message: {message[:50]}...") + if tool_output: + print(f"DEBUG cli_print_agent_messages: Received tool_output: {tool_output[:50]}...") + # Use the model from environment variable if available model_override = os.getenv('CAI_MODEL') if model_override: @@ -430,15 +803,29 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli timestamp = datetime.now().strftime("%H:%M:%S") - # Create a more hacker-like header + # Create header text = Text() + + # Check if the message has tool calls + has_tool_calls = False + if hasattr(message, 'tool_calls') and message.tool_calls: + has_tool_calls = True + elif isinstance(message, dict) and 'tool_calls' in message and message['tool_calls']: + has_tool_calls = True + + # Parse the message based on whether it has tool calls + if has_tool_calls: + parsed_message, tool_panels = parse_message_tool_call(message, tool_output) + else: + parsed_message = parse_message_content(message) + tool_panels = [] # Special handling for Reasoner Agent if agent_name == "Reasoner Agent": text.append(f"[{counter}] ", style="bold red") text.append(f"Agent: {agent_name} ", style="bold yellow") - if message: - text.append(f">> {message} ", style="green") + if parsed_message: + text.append(f">> {parsed_message} ", style="green") text.append(f"[{timestamp}", style="dim") if model: text.append(f" ({os.getenv('CAI_SUPPORT_MODEL')})", @@ -447,8 +834,8 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli else: text.append(f"[{counter}] ", style="bold cyan") text.append(f"Agent: {agent_name} ", style="bold green") - if message: - text.append(f">> {message} ", style="yellow") + if parsed_message: + text.append(f">> {parsed_message} ", style="yellow") text.append(f"[{timestamp}", style="dim") if model: text.append(f" ({model})", style="bold magenta") @@ -476,7 +863,6 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli ) text.append(tokens_text) - # Create a panel for better visual separation panel = Panel( text, border_style="red" if agent_name == "Reasoner Agent" else "blue", @@ -489,20 +875,27 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli ) console.print("\n") console.print(panel) + + # If there are tool panels, print them after the main message panel + # But only in non-streaming mode to avoid duplicates + is_streaming_enabled = os.getenv('CAI_STREAM', 'false').lower() == 'true' + if tool_panels and not is_streaming_enabled: + for tool_panel in tool_panels: + console.print(tool_panel) def create_agent_streaming_context(agent_name, counter, model): """Create a streaming context object that maintains state for streaming agent output.""" from rich.live import Live import shutil - # Use the model from environment variable if available + # Use the model from env if available model_override = os.getenv('CAI_MODEL') if model_override: model = model_override timestamp = datetime.now().strftime("%H:%M:%S") - # Determine terminal size for best display + # Terminal size for better display terminal_width, _ = shutil.get_terminal_size((100, 24)) panel_width = min(terminal_width - 4, 120) # Keep some margin @@ -527,18 +920,16 @@ def create_agent_streaming_context(agent_name, counter, model): Text.assemble(header, content, footer), border_style="blue", box=ROUNDED, - padding=(1, 2), # Add more padding for better readability + padding=(1, 2), title="[bold]Agent Streaming Response[/bold]", title_align="left", width=panel_width, - expand=True # Allow panel to expand to terminal width + expand=True ) - # Start the live display with a higher refresh rate - live = Live(panel, refresh_per_second=20, console=console) - live.start() + # Create Live display object but don't start it until we have content + live = Live(panel, refresh_per_second=20, console=console, auto_refresh=False) - # Return context object with all the elements needed for updating return { "live": live, "panel": panel, @@ -548,37 +939,57 @@ def create_agent_streaming_context(agent_name, counter, model): "timestamp": timestamp, "model": model, "agent_name": agent_name, - "panel_width": panel_width + "panel_width": panel_width, + "is_started": False # Track if we've started the display } def update_agent_streaming_content(context, text_delta): """Update the streaming content with new text.""" - # Add the new text to the content - context["content"].append(text_delta) + # Parse the text_delta to get just the content if needed + parsed_delta = parse_message_content(text_delta) + + # Skip empty updates to avoid showing an empty panel + if not parsed_delta or parsed_delta.strip() == "": + return + + # Add the parsed text to the content + context["content"].append(parsed_delta) # Update the live display with the latest content updated_panel = Panel( Text.assemble(context["header"], context["content"], context["footer"]), border_style="blue", box=ROUNDED, - padding=(1, 2), # Match padding from creation + padding=(1, 2), title="[bold]Agent Streaming Response[/bold]", title_align="left", width=context.get("panel_width", 100), - expand=True # Allow panel to expand to terminal width + expand=True ) + # Check if we need to start the display + if not context.get("is_started", False): + context["live"].start() + context["is_started"] = True + # Force an update with the new panel context["live"].update(updated_panel) context["panel"] = updated_panel def finish_agent_streaming(context, final_stats=None): """Finish the streaming session and display final stats if available.""" + # Check if there's actual content to display - don't show empty panels + if not context["content"] or context["content"].plain == "": + # If the display was never started, nothing to do + if not context.get("is_started", False): + return + # Otherwise, stop the display without showing final panel + context["live"].stop() + return + # If we have token stats, add them tokens_text = None if final_stats: - #print(f"\nDEBUG finish_agent_streaming: Received final_stats: {final_stats}") - interaction_input_tokens = final_stats.get("interaction_input_tokens") interaction_output_tokens = final_stats.get("interaction_output_tokens") interaction_reasoning_tokens = final_stats.get("interaction_reasoning_tokens") @@ -590,6 +1001,11 @@ def finish_agent_streaming(context, final_stats=None): interaction_cost = float(final_stats.get("interaction_cost", 0.0)) total_cost = float(final_stats.get("total_cost", 0.0)) + model_name = context.get("model", "") + # If model is not a string, use env + if not isinstance(model_name, str): + model_name = os.environ.get('CAI_MODEL', 'gpt-4o-mini') + if (interaction_input_tokens is not None and interaction_output_tokens is not None and interaction_reasoning_tokens is not None and @@ -599,9 +1015,9 @@ def finish_agent_streaming(context, final_stats=None): # Only calculate costs if they weren't provided or are zero if interaction_cost is None or interaction_cost == 0.0: - interaction_cost = calculate_model_cost(context["model"], interaction_input_tokens, interaction_output_tokens) + interaction_cost = calculate_model_cost(model_name, interaction_input_tokens, interaction_output_tokens) if total_cost is None or total_cost == 0.0: - total_cost = calculate_model_cost(context["model"], total_input_tokens, total_output_tokens) + total_cost = calculate_model_cost(model_name, total_input_tokens, total_output_tokens) tokens_text = _create_token_display( interaction_input_tokens, @@ -610,12 +1026,11 @@ def finish_agent_streaming(context, final_stats=None): total_input_tokens, total_output_tokens, total_reasoning_tokens, - context["model"], + model_name, # string model name! interaction_cost, total_cost ) - # Create the final panel with stats final_panel = Panel( Text.assemble( context["header"], @@ -626,7 +1041,7 @@ def finish_agent_streaming(context, final_stats=None): ), border_style="blue", box=ROUNDED, - padding=(1, 2), # Match padding from creation + padding=(1, 2), title="[bold]Agent Streaming Response[/bold]", title_align="left", width=context.get("panel_width", 100), @@ -643,42 +1058,211 @@ def finish_agent_streaming(context, final_stats=None): # Stop the live display context["live"].stop() -def calculate_model_cost(model_name, input_tokens, output_tokens): +def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execution_info=None, token_info=None): """ - Calculate the cost for a given model based on token usage. + Print a tool call output to the command line. + Similar to cli_print_tool_call but for the output of the tool. Args: - model_name: The name of the model being used - input_tokens: Number of input tokens used - output_tokens: Number of output tokens used - - Returns: - float: The calculated cost in dollars + tool_name: Name of the tool + args: Arguments passed to the tool + output: The output of the tool + call_id: Optional call ID for streaming updates + execution_info: Optional execution information + token_info: Optional token information """ - # Fetch model pricing data from LiteLLM GitHub repository - LITELLM_URL = ( - "https://raw.githubusercontent.com/BerriAI/litellm/main/" - "model_prices_and_context_window.json" - ) + # 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' + if is_streaming_enabled and not call_id: + # Skip all non-streaming tool output in streaming mode + return + + # Track seen call IDs to prevent duplicate panels + if not hasattr(cli_print_tool_output, '_seen_calls'): + cli_print_tool_output._seen_calls = {} + + # For streaming updates, only show updates for the same call_id + # but allow the first appearance of each call_id + if call_id: + call_key = f"{call_id}:{output[:20]}" # Use first 20 chars as fingerprint with call_id + + # Skip if we've seen this exact output for this call_id before + if call_key in cli_print_tool_output._seen_calls: + return + + # Mark as seen + cli_print_tool_output._seen_calls[call_key] = True + + # Limit cache size to prevent memory growth + if len(cli_print_tool_output._seen_calls) > 1000: + # Keep only the most recent 500 entries + cli_print_tool_output._seen_calls = { + k: cli_print_tool_output._seen_calls[k] + for k in list(cli_print_tool_output._seen_calls.keys())[-500:] + } + + # Try to use Rich for better formatting if available try: - import requests - response = requests.get(LITELLM_URL, timeout=2) - if response.status_code == 200: - model_pricing_data = response.json() + from rich.console import Console + from rich.panel import Panel + from rich.text import Text + from rich.box import ROUNDED + + # Create a console for output + console = Console() + + # Format arguments as a string if they are a dictionary + if isinstance(args, dict): + args_str = ", ".join([f"{key}='{value}'" for key, value in args.items()]) + else: + args_str = str(args) + + # 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") - # Get pricing info for the model - pricing_info = model_pricing_data.get(model_name, {}) - input_cost_per_token = pricing_info.get("input_cost_per_token", 0) - output_cost_per_token = pricing_info.get("output_cost_per_token", 0) + # Create content text with the output + content = Text(output) - # Calculate costs - input_cost = input_tokens * input_cost_per_token - output_cost = output_tokens * output_cost_per_token + # Create the panel for display + 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]", + border_style="blue", + padding=(1, 2), + box=ROUNDED, + title_align="left" + ) - return input_cost + output_cost - except Exception: - # If we can't fetch pricing data, return 0 - pass - - return 0.0 \ No newline at end of file + # Display using Rich + 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 + 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) + + 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) + + # Add the main 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) + + # Create the panel with all sections + panel = Panel( + Text.assemble(*sections), + title=f"[bold green]Tool Output: {tool_call}[/bold green]", + subtitle=subtitle, + border_style="green", + padding=(1, 2), + box=ROUNDED + ) + + # Display the panel + console.print(panel) + + except ImportError: + # Fall back to simple formatting if Rich is not available + # Format arguments as a string if they are a dictionary + if isinstance(args, dict): + args_str = ", ".join([f"{key}='{value}'" for key, value in args.items()]) + 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 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 + 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]" + else: + tool_call += 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) + + 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")) + + # Print the actual output + print(output) + print() \ No newline at end of file