From a47b61b96b49fd0765479e4185ca555cc2c13285 Mon Sep 17 00:00:00 2001 From: lidia9 Date: Tue, 22 Apr 2025 12:59:42 +0000 Subject: [PATCH 1/9] FIX UI - CAI_STREAM=false - print messages well parsed & parsed tool calls --- src/cai/cli.py | 2 +- src/cai/util.py | 195 ++++++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 189 insertions(+), 8 deletions(-) diff --git a/src/cai/cli.py b/src/cai/cli.py index 677914db..c67ec1f1 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -355,7 +355,7 @@ 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 diff --git a/src/cai/util.py b/src/cai/util.py index 4bab6da4..49ccce28 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -413,6 +413,165 @@ 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 content. + Sample of message object: + Message( + content='Hello! How can I assist you today?', + role='assistant', + tool_calls=None, + function_call=None, + provider_specific_fields={'refusal': None}, + annotations=[] + ) + + 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 the message has already been displayed, return empty string to avoid duplication + # This is a hacky approach but should work + + # 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=command, args=args) + and optionally shows the tool output in a separate 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 first (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 + + for tool_call in tool_calls: + # Extract tool name and arguments + tool_name = None + args_dict = {} + + # 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 we have a valid name + if tool_name: + # Format in the style shown in screenshot: tool_name(command=command, args=args) + tool_text = Text() + + # Start with the tool name in green + tool_text.append(f"{tool_name}", style="green") + + # Create the arguments list in the format (key=value, key=value) + args_parts = [] + for key, value in args_dict.items(): + # Format based on value type + if isinstance(value, bool): + args_parts.append(f"{key}={value}") + elif value == "" or value is None: + args_parts.append(f"{key}=") + else: + # If the value contains spaces or special chars, wrap it in quotes + if isinstance(value, str) and (' ' in value or '/' in value): + args_parts.append(f'{key}="{value}"') + else: + args_parts.append(f"{key}={value}") + + # Add the arguments in parentheses after the tool name + if args_parts: + tool_text.append("(", style="yellow") + tool_text.append(", ".join(args_parts), style="yellow") + tool_text.append(")", style="yellow") + + # Create the tool call panel (blue border) + tool_panel = Panel( + tool_text, + border_style="blue", + box=ROUNDED, + padding=(1, 2), + title="[bold]Tool Execution[/bold]", + title_align="left", + expand=True + ) + + tool_panels.append(tool_panel) + + # If there's a tool output, create a separate panel for it + if tool_output and tool_output.strip(): + output_panel = Panel( + Text(tool_output, style="yellow"), + border_style="red", + box=ROUNDED, + padding=(1, 2), + title="[bold]Tool Output[/bold]", + title_align="left", + expand=True + ) + tool_panels.append(output_panel) + + return content, tool_panels + 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,7 +580,8 @@ 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.""" # Use the model from environment variable if available model_override = os.getenv('CAI_MODEL') @@ -432,13 +592,27 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli # Create a more hacker-like 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 +621,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") @@ -489,6 +663,10 @@ 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 + 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.""" @@ -553,8 +731,11 @@ def create_agent_streaming_context(agent_name, counter, model): 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) + + # Add the parsed text to the content + context["content"].append(parsed_delta) # Update the live display with the latest content updated_panel = Panel( From 09989e9c1022eb131f7a0184b7adad4f81677fee Mon Sep 17 00:00:00 2001 From: lidia9 Date: Tue, 22 Apr 2025 15:12:49 +0000 Subject: [PATCH 2/9] FIX UI - CAI_STREAM=false - one panel for each message, tool call, output --- .../agents/models/openai_chatcompletions.py | 173 ++++++++++-- src/cai/util.py | 255 +++++++++++++++--- 2 files changed, 371 insertions(+), 57 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 4031a7af..f658e183 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,80 @@ 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 + + # 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, + execution_info=execution_info, + token_info=token_info + ) + + # Continue with normal processing flush_assistant_message() msg: ChatCompletionToolMessageParam = { "role": "tool", diff --git a/src/cai/util.py b/src/cai/util.py index 49ccce28..d6beca85 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -415,16 +415,8 @@ def _create_token_display( # pylint: disable=too-many-arguments,too-many-locals def parse_message_content(message): """ - Parse a message object to extract its content. - Sample of message object: - Message( - content='Hello! How can I assist you today?', - role='assistant', - tool_calls=None, - function_call=None, - provider_specific_fields={'refusal': None}, - annotations=[] - ) + 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 @@ -455,7 +447,7 @@ 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=command, args=args) - and optionally shows the tool output in a separate panel. + and shows the tool output in the same panel. Args: message: A Message object or dict with content and tool_calls attributes @@ -468,6 +460,10 @@ def parse_message_tool_call(message, tool_output=None): content = "" tool_panels = [] + # Debug the incoming tool_output + if tool_output: + print(f"DEBUG parse_message_tool_call: Received tool_output: {tool_output[:50]}...") + # Extract the content text first (LLM's inference) if isinstance(message, str): content = message @@ -488,11 +484,21 @@ def parse_message_tool_call(message, tool_output=None): 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 + + # Extract call_id for debugging + if hasattr(tool_call, 'id'): + call_id = tool_call.id + elif isinstance(tool_call, dict) and 'id' in tool_call: + call_id = tool_call['id'] + + print(f"DEBUG parse_message_tool_call: Processing tool_call with call_id={call_id}") # Handle different formats of tool_call objects if hasattr(tool_call, 'function'): @@ -517,36 +523,48 @@ def parse_message_tool_call(message, tool_output=None): # Create a panel for this tool call if we have a valid name if tool_name: - # Format in the style shown in screenshot: tool_name(command=command, args=args) + # Create content for the panel + panel_content = [] + + # Start with the tool name and arguments tool_text = Text() + tool_text.append(f"{tool_name}", style="bold #00BCD4") # Cyan (timestamp color from theme) in bold - # Start with the tool name in green - tool_text.append(f"{tool_name}", style="green") - - # Create the arguments list in the format (key=value, key=value) + # Format arguments args_parts = [] for key, value in args_dict.items(): - # Format based on value type if isinstance(value, bool): args_parts.append(f"{key}={value}") elif value == "" or value is None: args_parts.append(f"{key}=") else: - # If the value contains spaces or special chars, wrap it in quotes if isinstance(value, str) and (' ' in value or '/' in value): args_parts.append(f'{key}="{value}"') else: args_parts.append(f"{key}={value}") - # Add the arguments in parentheses after the tool name if args_parts: tool_text.append("(", style="yellow") tool_text.append(", ".join(args_parts), style="yellow") tool_text.append(")", style="yellow") - # Create the tool call panel (blue border) + panel_content.append(tool_text) + + # Add tool output to the same panel if available + if tool_output: + print(f"DEBUG parse_message_tool_call: Adding tool_output to panel: {tool_output[:50]}...") + divider_text = Text("\n" + "─" * 50, style="dim") + output_text = Text("\nOutput:", style="bold #C0C0C0") # Change to silver/gray + output_text.append(f"\n{tool_output}", style="#C0C0C0") # Change to silver/gray + + panel_content.append(divider_text) + panel_content.append(output_text) + else: + print("DEBUG parse_message_tool_call: No tool_output available to add to panel") + + # Create a single panel with both tool call and output tool_panel = Panel( - tool_text, + Group(*panel_content), border_style="blue", box=ROUNDED, padding=(1, 2), @@ -556,22 +574,17 @@ def parse_message_tool_call(message, tool_output=None): ) tool_panels.append(tool_panel) - - # If there's a tool output, create a separate panel for it - if tool_output and tool_output.strip(): - output_panel = Panel( - Text(tool_output, style="yellow"), - border_style="red", - box=ROUNDED, - padding=(1, 2), - title="[bold]Tool Output[/bold]", - title_align="left", - expand=True - ) - tool_panels.append(output_panel) 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, @@ -583,6 +596,13 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli 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: @@ -862,4 +882,169 @@ def calculate_model_cost(model_name, input_tokens, output_tokens): # If we can't fetch pricing data, return 0 pass - return 0.0 \ No newline at end of file + return 0.0 + +def cli_print_tool_output(tool_name, args, output, call_id=None, execution_info=None, token_info=None): + """ + Print tool execution and output in a single unified panel. + + Args: + tool_name: Name of the tool that was executed + args: Arguments passed to the tool + output: Output from the tool execution + call_id: Optional ID of the tool call + execution_info: Dictionary with execution timing information + token_info: Dictionary with token usage information + """ + from rich.panel import Panel + from rich.text import Text + from rich.box import ROUNDED + from rich.console import Group + + # Track which tool outputs we've already printed to avoid duplicates + # Use a module-level dictionary to track call_ids we've seen + if not hasattr(cli_print_tool_output, '_seen_calls'): + cli_print_tool_output._seen_calls = {} + + # If we have a call_id, check if we've already printed this output + # If no call_id, use a combination of tool_name and args as a key + call_key = call_id if call_id else f"{tool_name}:{args}" + + # If we've already seen this output, don't print it again + if call_key in cli_print_tool_output._seen_calls: + return + + # Mark this call as seen to avoid duplicates + cli_print_tool_output._seen_calls[call_key] = True + + # Get execution time if available, otherwise don't show it + execution_time = None + if execution_info: + # Format execution time if available + total_time = execution_info.get('total_time', 0) + tool_time = execution_info.get('tool_time', 0) + if total_time > 0: + total_time_str = f"{int(total_time // 60)}m {total_time % 60:.1f}s" + tool_time_str = f"{tool_time:.1f}s" + execution_time = f"Total: {total_time_str} | Tool: {tool_time_str}" + + # Format the tool and arguments in the first panel + if isinstance(args, dict): + # Format as key=value pairs + args_str = ", ".join(f"{k}={v}" for k, v in args.items()) + display_str = f"{tool_name}({args_str})" + else: + # If args is just a string + display_str = f"{tool_name}({args})" + + # Create content for the first panel - just tool name and args + tool_text = Text() + tool_text.append(f"{tool_name}", style="#00BCD4") # Cyan (timestamp color from theme) + + # Add the arguments in their original color + if isinstance(args, dict): + args_str = ", ".join(f"{k}={v}" for k, v in args.items()) + tool_text.append("(", style="yellow") + tool_text.append(args_str, style="yellow") + tool_text.append(")", style="yellow") + else: + tool_text.append("(", style="yellow") + tool_text.append(f"{args}", style="yellow") + tool_text.append(")", style="yellow") + + # Create the first panel - just shows the tool name and args + first_panel = Panel( + tool_text, + border_style="blue", + box=ROUNDED, + padding=(1, 2), + title="[bold]Tool Execution[/bold]", + title_align="left", + expand=True + ) + console.print(first_panel) + + # Create content for the second panel + panel_content = [] + + # For the second panel, only show execution time without repeating the tool name + if execution_time: + exec_text = Text() + exec_text.append("[", style="dim") + exec_text.append(f"{execution_time}", style="magenta") + exec_text.append("]", style="dim") + panel_content.append(exec_text) + + # Add the tool output - change from yellow to silver/gray + if output and output.strip(): + output_text = Text("\n" if execution_time else "") + # Use a silver/gray color (#C0C0C0) for the output + output_text.append(output.strip(), style="#C0C0C0") + panel_content.append(output_text) + + # Add token display if token info is available + 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) + + # Calculate costs if possible + interaction_cost = calculate_model_cost(model, interaction_input_tokens, interaction_output_tokens) + total_cost = calculate_model_cost(model, total_input_tokens, total_output_tokens) + + # Calculate context usage + context_pct = 0 + if model: + context_pct = interaction_input_tokens / get_model_input_tokens(model) * 100 + + # Create token display + tokens_text = Text("\n" if panel_content else "") + tokens_text.append('(tokens) Interaction: ', style="dim") + 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") + tokens_text.append(f'(${interaction_cost:.4f}) ', style="bold") + tokens_text.append('| ', style="dim") + tokens_text.append('Total: ', style="dim") + 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") + tokens_text.append(f'(${total_cost:.4f}) ', style="bold") + tokens_text.append('| ', style="dim") + tokens_text.append(f'Context: {context_pct:.1f}% ', style="bold") + + # Context indicator + if context_pct < 50: + indicator = "🟩" + color_local = "green" + elif context_pct < 80: + indicator = "🟨" + color_local = "yellow" + else: + indicator = "🟥" + color_local = "red" + + tokens_text.append(f"{indicator}", style=color_local) + + # Add max context window size + if model: + max_tokens = get_model_input_tokens(model) + tokens_text.append(f" ({max_tokens})", style="dim") + + panel_content.append(tokens_text) + + # Only create and print the second panel if we have content for it + if panel_content: + # Create the second panel - shows execution time, output, and token info + second_panel = Panel( + Group(*panel_content), + border_style="blue", + box=ROUNDED, + padding=(1, 2), + expand=True + ) + console.print(second_panel) \ No newline at end of file From c50ff9221b70f2d4b23d025ced3c6e48536b9fbe Mon Sep 17 00:00:00 2001 From: lidia9 Date: Tue, 22 Apr 2025 15:17:14 +0000 Subject: [PATCH 3/9] FIX UI - remove duplicated tool call print --- src/cai/util.py | 118 +++++++++++++++++++----------------------------- 1 file changed, 47 insertions(+), 71 deletions(-) diff --git a/src/cai/util.py b/src/cai/util.py index d6beca85..13e4ff7a 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -522,58 +522,38 @@ def parse_message_tool_call(message, tool_output=None): args_dict = {"raw_arguments": tool_call['function']['arguments']} # Create a panel for this tool call if we have a valid name - if tool_name: - # Create content for the panel + # Don't show the tool execution panel here - that will be handled by cli_print_tool_output + # We'll only 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 = [] - # Start with the tool name and arguments - tool_text = Text() - tool_text.append(f"{tool_name}", style="bold #00BCD4") # Cyan (timestamp color from theme) in bold + # 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 - # Format arguments - args_parts = [] - for key, value in args_dict.items(): - if isinstance(value, bool): - args_parts.append(f"{key}={value}") - elif value == "" or value is None: - args_parts.append(f"{key}=") - else: - if isinstance(value, str) and (' ' in value or '/' in value): - args_parts.append(f'{key}="{value}"') - else: - args_parts.append(f"{key}={value}") + panel_content.append(output_text) - if args_parts: - tool_text.append("(", style="yellow") - tool_text.append(", ".join(args_parts), style="yellow") - tool_text.append(")", style="yellow") - - panel_content.append(tool_text) - - # Add tool output to the same panel if available - if tool_output: - print(f"DEBUG parse_message_tool_call: Adding tool_output to panel: {tool_output[:50]}...") - divider_text = Text("\n" + "─" * 50, style="dim") - output_text = Text("\nOutput:", style="bold #C0C0C0") # Change to silver/gray - output_text.append(f"\n{tool_output}", style="#C0C0C0") # Change to silver/gray - - panel_content.append(divider_text) - panel_content.append(output_text) - else: - print("DEBUG parse_message_tool_call: No tool_output available to add to panel") - - # Create a single panel with both tool call and output + # Create a panel with just the output tool_panel = Panel( Group(*panel_content), border_style="blue", box=ROUNDED, padding=(1, 2), - title="[bold]Tool Execution[/bold]", + 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 @@ -901,8 +881,7 @@ def cli_print_tool_output(tool_name, args, output, call_id=None, execution_info= from rich.box import ROUNDED from rich.console import Group - # Track which tool outputs we've already printed to avoid duplicates - # Use a module-level dictionary to track call_ids we've seen + # Track which tool calls we've already printed to avoid duplicates if not hasattr(cli_print_tool_output, '_seen_calls'): cli_print_tool_output._seen_calls = {} @@ -917,27 +896,15 @@ def cli_print_tool_output(tool_name, args, output, call_id=None, execution_info= # Mark this call as seen to avoid duplicates cli_print_tool_output._seen_calls[call_key] = True - # Get execution time if available, otherwise don't show it - execution_time = None - if execution_info: - # Format execution time if available - total_time = execution_info.get('total_time', 0) - tool_time = execution_info.get('tool_time', 0) - if total_time > 0: - total_time_str = f"{int(total_time // 60)}m {total_time % 60:.1f}s" - tool_time_str = f"{tool_time:.1f}s" - execution_time = f"Total: {total_time_str} | Tool: {tool_time_str}" - - # Format the tool and arguments in the first panel + # Format the tool and arguments if isinstance(args, dict): # Format as key=value pairs args_str = ", ".join(f"{k}={v}" for k, v in args.items()) - display_str = f"{tool_name}({args_str})" else: # If args is just a string - display_str = f"{tool_name}({args})" + args_str = args - # Create content for the first panel - just tool name and args + # Create content for the tool execution panel tool_text = Text() tool_text.append(f"{tool_name}", style="#00BCD4") # Cyan (timestamp color from theme) @@ -952,8 +919,8 @@ def cli_print_tool_output(tool_name, args, output, call_id=None, execution_info= tool_text.append(f"{args}", style="yellow") tool_text.append(")", style="yellow") - # Create the first panel - just shows the tool name and args - first_panel = Panel( + # Create the tool execution panel + tool_panel = Panel( tool_text, border_style="blue", box=ROUNDED, @@ -962,22 +929,30 @@ def cli_print_tool_output(tool_name, args, output, call_id=None, execution_info= title_align="left", expand=True ) - console.print(first_panel) + console.print(tool_panel) - # Create content for the second panel + # Create content for the output panel panel_content = [] - # For the second panel, only show execution time without repeating the tool name - if execution_time: - exec_text = Text() - exec_text.append("[", style="dim") - exec_text.append(f"{execution_time}", style="magenta") - exec_text.append("]", style="dim") - panel_content.append(exec_text) + # Add execution time if available + if execution_info: + # Format execution time if available + total_time = execution_info.get('total_time', 0) + tool_time = execution_info.get('tool_time', 0) + if total_time > 0: + total_time_str = f"{int(total_time // 60)}m {total_time % 60:.1f}s" + tool_time_str = f"{tool_time:.1f}s" + execution_time = f"Total: {total_time_str} | Tool: {tool_time_str}" + + exec_text = Text() + exec_text.append("[", style="dim") + exec_text.append(f"{execution_time}", style="magenta") + exec_text.append("]", style="dim") + panel_content.append(exec_text) # Add the tool output - change from yellow to silver/gray if output and output.strip(): - output_text = Text("\n" if execution_time else "") + output_text = Text("\n" if panel_content else "") # Use a silver/gray color (#C0C0C0) for the output output_text.append(output.strip(), style="#C0C0C0") panel_content.append(output_text) @@ -1037,14 +1012,15 @@ def cli_print_tool_output(tool_name, args, output, call_id=None, execution_info= panel_content.append(tokens_text) - # Only create and print the second panel if we have content for it + # Only create and print the output panel if we have content for it if panel_content: - # Create the second panel - shows execution time, output, and token info - second_panel = Panel( + # Create the output panel + output_panel = Panel( Group(*panel_content), border_style="blue", box=ROUNDED, padding=(1, 2), + title="[bold]Tool Output[/bold]", # Changed title to clarify this is output only expand=True ) - console.print(second_panel) \ No newline at end of file + console.print(output_panel) \ No newline at end of file From 70b7efd5e39219f79254c86722cfb9f7c0a67720 Mon Sep 17 00:00:00 2001 From: lidia9 Date: Wed, 23 Apr 2025 09:58:29 +0000 Subject: [PATCH 4/9] UI - fix display metrics, token count, pricing, session cost and so on --- src/cai/util.py | 510 +++++++++++++++++++++++++++++++----------------- 1 file changed, 333 insertions(+), 177 deletions(-) diff --git a/src/cai/util.py b/src/cai/util.py index 13e4ff7a..504ccfea 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -16,6 +16,213 @@ 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 + +# Near the top of the file with other global variables +_LAST_CALCULATION_ID = "" # Track the last calculation to avoid duplicates + +# Global variables to track costs +_CAI_SESSION_TOTAL_COST = 0.0 # Tracks cost across entire CAI session +_CURRENT_AGENT_TOTAL_COST = 0.0 # Tracks total cost for current agent run + +# Add this near the top of the file to keep track of the last total cost +_LAST_TOTAL_COST = 0.0 + +# Near the top of the file with other global variables +_MODEL_PRICING_CACHE = {} # Cache for model pricing data to avoid redundant lookups + +# 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 + # print(f"SESSION UPDATE: Adding ${new_cost:.6f} to session total (${old_total:.6f} → ${self.session_total_cost:.6f})") + + 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 + + # Log calculation with optional label + # log_prefix = f"{label}: " if label else "COST CALCULATION: " + # print(f"{log_prefix}model={model_name}") + # print(f" • Input: {input_tokens} tokens × ${input_cost_per_token:.8f} = ${input_cost:.6f}") + # print(f" • Output: {output_tokens} tokens × ${output_cost_per_token:.8f} = ${output_cost:.6f}") + # print(f" • Total Cost: ${total_cost:.6f}") + + # 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") + + # Debug: track the difference from last interaction + # cost_diff = self.interaction_cost - self.last_interaction_cost + # print(f"DEBUG: Interaction cost changed by ${cost_diff:.6f} (${self.last_interaction_cost:.6f} → ${self.interaction_cost:.6f})") + 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 + + # SIMPLIFIED APPROACH: Get previous total and add current interaction cost + previous_total = self.current_agent_total_cost + + # Instead of recalculating the total, simply add the new interaction cost + # This is the key change to ensure consistency + 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 + + # print(f"DEBUG: New total calculated by adding interaction cost: ${previous_total:.6f} + ${cost_diff:.6f} = ${new_total_cost:.6f}") + + # 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 @@ -331,9 +538,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 +601,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") @@ -424,10 +683,7 @@ def parse_message_content(message): Returns: str: The extracted content as a string """ - # Check if this is a duplicate print from OpenAIChatCompletionsModel - # If the message has already been displayed, return empty string to avoid duplication - # This is a hacky approach but should work - + # Check if this is a duplicate print from OpenAIChatCompletionsModel # If message is already a string, return it if isinstance(message, str): return message @@ -446,8 +702,8 @@ def parse_message_content(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=command, args=args) - and shows the tool output in the same panel. + 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 @@ -461,10 +717,10 @@ def parse_message_tool_call(message, tool_output=None): tool_panels = [] # Debug the incoming tool_output - if tool_output: - print(f"DEBUG parse_message_tool_call: Received tool_output: {tool_output[:50]}...") + #if tool_output: + # print(f"DEBUG parse_message_tool_call: Received tool_output: {tool_output[:50]}...") - # Extract the content text first (LLM's inference) + # Extract the content text (LLM's inference) if isinstance(message, str): content = message elif hasattr(message, 'content') and message.content is not None: @@ -492,13 +748,13 @@ def parse_message_tool_call(message, tool_output=None): args_dict = {} call_id = None - # Extract call_id for debugging - if hasattr(tool_call, 'id'): - call_id = tool_call.id - elif isinstance(tool_call, dict) and 'id' in tool_call: - call_id = tool_call['id'] + ## Extract call_id for debugging + #if hasattr(tool_call, 'id'): + # call_id = tool_call.id + #elif isinstance(tool_call, dict) and 'id' in tool_call: + # call_id = tool_call['id'] - print(f"DEBUG parse_message_tool_call: Processing tool_call with call_id={call_id}") + #print(f"DEBUG parse_message_tool_call: Processing tool_call with call_id={call_id}") # Handle different formats of tool_call objects if hasattr(tool_call, 'function'): @@ -521,9 +777,9 @@ def parse_message_tool_call(message, tool_output=None): except: args_dict = {"raw_arguments": tool_call['function']['arguments']} - # Create a panel for this tool call if we have a valid name - # Don't show the tool execution panel here - that will be handled by cli_print_tool_output - # We'll only pass on tool info to generate panels for display in cli_print_agent_messages + # 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 = [] @@ -590,7 +846,7 @@ 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 @@ -650,7 +906,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", @@ -673,14 +928,14 @@ def create_agent_streaming_context(agent_name, counter, model): 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 @@ -705,18 +960,17 @@ 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 + # Start the live display live = Live(panel, refresh_per_second=20, console=console) live.start() - # Return context object with all the elements needed for updating return { "live": live, "panel": panel, @@ -742,11 +996,11 @@ def update_agent_streaming_content(context, text_delta): 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 ) # Force an update with the new panel @@ -758,8 +1012,6 @@ def finish_agent_streaming(context, final_stats=None): # 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") @@ -771,6 +1023,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 @@ -780,9 +1037,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, @@ -791,12 +1048,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"], @@ -807,7 +1063,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), @@ -824,46 +1080,6 @@ 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): - """ - Calculate the cost for a given model based on token usage. - - 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 - """ - # Fetch model pricing data from LiteLLM GitHub repository - 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) - - # Calculate costs - input_cost = input_tokens * input_cost_per_token - output_cost = output_tokens * output_cost_per_token - - return input_cost + output_cost - except Exception: - # If we can't fetch pricing data, return 0 - pass - - return 0.0 - def cli_print_tool_output(tool_name, args, output, call_id=None, execution_info=None, token_info=None): """ Print tool execution and output in a single unified panel. @@ -874,7 +1090,7 @@ def cli_print_tool_output(tool_name, args, output, call_id=None, execution_info= output: Output from the tool execution call_id: Optional ID of the tool call execution_info: Dictionary with execution timing information - token_info: Dictionary with token usage information + token_info: Dictionary with token usage information (ignored - no cost display in tool panels) """ from rich.panel import Panel from rich.text import Text @@ -889,7 +1105,7 @@ def cli_print_tool_output(tool_name, args, output, call_id=None, execution_info= # If no call_id, use a combination of tool_name and args as a key call_key = call_id if call_id else f"{tool_name}:{args}" - # If we've already seen this output, don't print it again + # Avoid printing duplicates if call_key in cli_print_tool_output._seen_calls: return @@ -897,18 +1113,16 @@ def cli_print_tool_output(tool_name, args, output, call_id=None, execution_info= cli_print_tool_output._seen_calls[call_key] = True # Format the tool and arguments - if isinstance(args, dict): - # Format as key=value pairs + if isinstance(args, dict): #key=value pairs args_str = ", ".join(f"{k}={v}" for k, v in args.items()) - else: - # If args is just a string + else: #single string args_str = args # Create content for the tool execution panel tool_text = Text() - tool_text.append(f"{tool_name}", style="#00BCD4") # Cyan (timestamp color from theme) + tool_text.append(f"{tool_name}", style="#00BCD4") - # Add the arguments in their original color + # Add the arguments if isinstance(args, dict): args_str = ", ".join(f"{k}={v}" for k, v in args.items()) tool_text.append("(", style="yellow") @@ -936,7 +1150,6 @@ def cli_print_tool_output(tool_name, args, output, call_id=None, execution_info= # Add execution time if available if execution_info: - # Format execution time if available total_time = execution_info.get('total_time', 0) tool_time = execution_info.get('tool_time', 0) if total_time > 0: @@ -950,77 +1163,20 @@ def cli_print_tool_output(tool_name, args, output, call_id=None, execution_info= exec_text.append("]", style="dim") panel_content.append(exec_text) - # Add the tool output - change from yellow to silver/gray + # Add the tool output if output and output.strip(): output_text = Text("\n" if panel_content else "") - # Use a silver/gray color (#C0C0C0) for the output output_text.append(output.strip(), style="#C0C0C0") panel_content.append(output_text) - # Add token display if token info is available - 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) - - # Calculate costs if possible - interaction_cost = calculate_model_cost(model, interaction_input_tokens, interaction_output_tokens) - total_cost = calculate_model_cost(model, total_input_tokens, total_output_tokens) - - # Calculate context usage - context_pct = 0 - if model: - context_pct = interaction_input_tokens / get_model_input_tokens(model) * 100 - - # Create token display - tokens_text = Text("\n" if panel_content else "") - tokens_text.append('(tokens) Interaction: ', style="dim") - 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") - tokens_text.append(f'(${interaction_cost:.4f}) ', style="bold") - tokens_text.append('| ', style="dim") - tokens_text.append('Total: ', style="dim") - 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") - tokens_text.append(f'(${total_cost:.4f}) ', style="bold") - tokens_text.append('| ', style="dim") - tokens_text.append(f'Context: {context_pct:.1f}% ', style="bold") - - # Context indicator - if context_pct < 50: - indicator = "🟩" - color_local = "green" - elif context_pct < 80: - indicator = "🟨" - color_local = "yellow" - else: - indicator = "🟥" - color_local = "red" - - tokens_text.append(f"{indicator}", style=color_local) - - # Add max context window size - if model: - max_tokens = get_model_input_tokens(model) - tokens_text.append(f" ({max_tokens})", style="dim") - - panel_content.append(tokens_text) - - # Only create and print the output panel if we have content for it if panel_content: - # Create the output panel output_panel = Panel( Group(*panel_content), border_style="blue", box=ROUNDED, padding=(1, 2), - title="[bold]Tool Output[/bold]", # Changed title to clarify this is output only + title="[bold]Tool Output[/bold]", + title_align="left", expand=True ) console.print(output_panel) \ No newline at end of file From 2d8f5abb54f48faaccc14d58921aea3b78fd8e20 Mon Sep 17 00:00:00 2001 From: lidia9 Date: Wed, 23 Apr 2025 10:03:51 +0000 Subject: [PATCH 5/9] UI - fix display metrics, token count, pricing, session cost - REMOVE some prints --- src/cai/util.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/cai/util.py b/src/cai/util.py index 504ccfea..e0eee348 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -189,11 +189,10 @@ class CostTracker: self.current_agent_output_tokens = total_output_tokens self.current_agent_reasoning_tokens = total_reasoning_tokens - # SIMPLIFIED APPROACH: Get previous total and add current interaction cost + # Get previous total and add current interaction cost previous_total = self.current_agent_total_cost - # Instead of recalculating the total, simply add the new interaction cost - # This is the key change to ensure consistency + # 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) From 66fb3037fdfa3a1f8d9e6c9513b2ff4d29a5edcc Mon Sep 17 00:00:00 2001 From: lidia9 Date: Wed, 23 Apr 2025 12:28:47 +0000 Subject: [PATCH 6/9] remove prints --- src/cai/util.py | 80 ++++++++++++------------------------------------- 1 file changed, 19 insertions(+), 61 deletions(-) diff --git a/src/cai/util.py b/src/cai/util.py index e0eee348..ab7107db 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -20,18 +20,6 @@ import atexit from dataclasses import dataclass, field from typing import Dict, Optional -# Near the top of the file with other global variables -_LAST_CALCULATION_ID = "" # Track the last calculation to avoid duplicates - -# Global variables to track costs -_CAI_SESSION_TOTAL_COST = 0.0 # Tracks cost across entire CAI session -_CURRENT_AGENT_TOTAL_COST = 0.0 # Tracks total cost for current agent run - -# Add this near the top of the file to keep track of the last total cost -_LAST_TOTAL_COST = 0.0 - -# Near the top of the file with other global variables -_MODEL_PRICING_CACHE = {} # Cache for model pricing data to avoid redundant lookups # Shared stats tracking object to maintain consistent costs across calls @dataclass @@ -70,8 +58,7 @@ class CostTracker: """Add cost to session total and log the update""" old_total = self.session_total_cost self.session_total_cost += new_cost - # print(f"SESSION UPDATE: Adding ${new_cost:.6f} to session total (${old_total:.6f} → ${self.session_total_cost:.6f})") - + def log_final_cost(self) -> None: """Display final cost information at exit""" print(f"\nTotal CAI Session Cost: ${self.session_total_cost:.6f}") @@ -133,14 +120,7 @@ class CostTracker: input_cost = input_tokens * input_cost_per_token output_cost = output_tokens * output_cost_per_token total_cost = input_cost + output_cost - - # Log calculation with optional label - # log_prefix = f"{label}: " if label else "COST CALCULATION: " - # print(f"{log_prefix}model={model_name}") - # print(f" • Input: {input_tokens} tokens × ${input_cost_per_token:.8f} = ${input_cost:.6f}") - # print(f" • Output: {output_tokens} tokens × ${output_cost_per_token:.8f} = ${output_cost:.6f}") - # print(f" • Total Cost: ${total_cost:.6f}") - + # Cache the result with full precision self.calculated_costs_cache[cache_key] = total_cost @@ -168,9 +148,6 @@ class CostTracker: model_name, input_tokens, output_tokens, label="OFFICIAL CALCULATION: Interaction") - # Debug: track the difference from last interaction - # cost_diff = self.interaction_cost - self.last_interaction_cost - # print(f"DEBUG: Interaction cost changed by ${cost_diff:.6f} (${self.last_interaction_cost:.6f} → ${self.interaction_cost:.6f})") self.last_interaction_cost = self.interaction_cost return self.interaction_cost @@ -202,8 +179,6 @@ class CostTracker: # Simply add the current interaction cost to the previous total cost_diff = self.interaction_cost new_total_cost = previous_total + cost_diff - - # print(f"DEBUG: New total calculated by adding interaction cost: ${previous_total:.6f} + ${cost_diff:.6f} = ${new_total_cost:.6f}") # Only add to session total if there's genuinely new cost (and it's positive) if cost_diff > 0: @@ -222,32 +197,27 @@ 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) @@ -715,10 +685,6 @@ def parse_message_tool_call(message, tool_output=None): content = "" tool_panels = [] - # Debug the incoming tool_output - #if tool_output: - # print(f"DEBUG parse_message_tool_call: Received tool_output: {tool_output[:50]}...") - # Extract the content text (LLM's inference) if isinstance(message, str): content = message @@ -747,14 +713,6 @@ def parse_message_tool_call(message, tool_output=None): args_dict = {} call_id = None - ## Extract call_id for debugging - #if hasattr(tool_call, 'id'): - # call_id = tool_call.id - #elif isinstance(tool_call, dict) and 'id' in tool_call: - # call_id = tool_call['id'] - - #print(f"DEBUG parse_message_tool_call: Processing tool_call with call_id={call_id}") - # Handle different formats of tool_call objects if hasattr(tool_call, 'function'): if hasattr(tool_call.function, 'name'): From 49dbebc76d4dcd4a7a834abfa1138d7e342d9062 Mon Sep 17 00:00:00 2001 From: lidia9 Date: Fri, 25 Apr 2025 10:33:12 +0200 Subject: [PATCH 7/9] FIX streaming output tool - still some visualization duplicates --- src/cai/cli.py | 4 +- src/cai/tools/common.py | 200 +++++++++++- .../reconnaissance/generic_linux_command.py | 18 +- src/cai/util.py | 305 ++++++++++++------ 4 files changed, 422 insertions(+), 105 deletions(-) diff --git a/src/cai/cli.py b/src/cai/cli.py index c67ec1f1..1d2fed05 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 diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 91a9752a..ebf6821e 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( @@ -252,9 +257,184 @@ 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" + header = Text() + header.append(tool_name, style="#00BCD4") + header.append("(", style="yellow") + header.append(f"command='{command}'", 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 blue]Tool Execution[/bold blue]", + subtitle="[bold green]Live Output[/bold green]", + border_style="blue", + 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 blue]Tool Execution[/bold blue]", + subtitle="[bold green]Live Output[/bold green]", + border_style="blue", + 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 blue]Tool Execution[/bold blue]", + subtitle="[bold green]Live Output[/bold green]", + border_style="blue", + 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 blue]Tool Execution[/bold blue]", + subtitle=f"[bold green]{completion_status}[/bold green]", + border_style="blue", + 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 + tool_args = {"command": command} + + # 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 +444,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 +470,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..7fa156e7 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,19 @@ 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 + call_id = str(uuid.uuid4())[:8] if stream else None + # Run the command with the appropriate parameters return run_command(full_command, ctf=ctf, async_mode=async_mode, session_id=session_id, - timeout=timeout) + timeout=timeout, stream=stream, call_id=call_id) diff --git a/src/cai/util.py b/src/cai/util.py index ab7107db..6397a230 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -924,9 +924,8 @@ def create_agent_streaming_context(agent_name, counter, model): expand=True ) - # Start the live display - 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 { "live": live, @@ -937,7 +936,8 @@ 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): @@ -945,6 +945,10 @@ def update_agent_streaming_content(context, 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) @@ -960,12 +964,26 @@ def update_agent_streaming_content(context, text_delta): 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: @@ -1037,103 +1055,204 @@ def finish_agent_streaming(context, final_stats=None): # Stop the live display context["live"].stop() -def cli_print_tool_output(tool_name, args, output, call_id=None, execution_info=None, token_info=None): +def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execution_info=None, token_info=None): """ - Print tool execution and output in a single unified panel. + Print a tool call output to the command line. + Similar to cli_print_tool_call but for the output of the tool. Args: - tool_name: Name of the tool that was executed + tool_name: Name of the tool args: Arguments passed to the tool - output: Output from the tool execution - call_id: Optional ID of the tool call - execution_info: Dictionary with execution timing information - token_info: Dictionary with token usage information (ignored - no cost display in tool panels) + output: The output of the tool + call_id: Optional call ID for streaming updates + execution_info: Optional execution information + token_info: Optional token information """ - from rich.panel import Panel - from rich.text import Text - from rich.box import ROUNDED - from rich.console import Group - - # Track which tool calls we've already printed to avoid duplicates - if not hasattr(cli_print_tool_output, '_seen_calls'): - cli_print_tool_output._seen_calls = {} - - # If we have a call_id, check if we've already printed this output - # If no call_id, use a combination of tool_name and args as a key - call_key = call_id if call_id else f"{tool_name}:{args}" - - # Avoid printing duplicates - if call_key in cli_print_tool_output._seen_calls: + # If it's an empty output, don't print anything + if not output and not call_id: return - # Mark this call as seen to avoid duplicates - cli_print_tool_output._seen_calls[call_key] = True - - # Format the tool and arguments - if isinstance(args, dict): #key=value pairs - args_str = ", ".join(f"{k}={v}" for k, v in args.items()) - else: #single string - args_str = args - - # Create content for the tool execution panel - tool_text = Text() - tool_text.append(f"{tool_name}", style="#00BCD4") - - # Add the arguments - if isinstance(args, dict): - args_str = ", ".join(f"{k}={v}" for k, v in args.items()) - tool_text.append("(", style="yellow") - tool_text.append(args_str, style="yellow") - tool_text.append(")", style="yellow") - else: - tool_text.append("(", style="yellow") - tool_text.append(f"{args}", style="yellow") - tool_text.append(")", style="yellow") - - # Create the tool execution panel - tool_panel = Panel( - tool_text, - border_style="blue", - box=ROUNDED, - padding=(1, 2), - title="[bold]Tool Execution[/bold]", - title_align="left", - expand=True - ) - console.print(tool_panel) - - # Create content for the output panel - panel_content = [] - - # Add execution time if available - if execution_info: - total_time = execution_info.get('total_time', 0) - tool_time = execution_info.get('tool_time', 0) - if total_time > 0: - total_time_str = f"{int(total_time // 60)}m {total_time % 60:.1f}s" - tool_time_str = f"{tool_time:.1f}s" - execution_time = f"Total: {total_time_str} | Tool: {tool_time_str}" + # 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 - exec_text = Text() - exec_text.append("[", style="dim") - exec_text.append(f"{execution_time}", style="magenta") - exec_text.append("]", style="dim") - panel_content.append(exec_text) + # 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:] + } - # Add the tool output - if output and output.strip(): - output_text = Text("\n" if panel_content else "") - output_text.append(output.strip(), style="#C0C0C0") - panel_content.append(output_text) - - if panel_content: - output_panel = Panel( - Group(*panel_content), - border_style="blue", - box=ROUNDED, + # Try to use Rich for better formatting if available + try: + 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") + + # Create content text with the output + content = Text(output) + + # Create the panel for display + panel = Panel( + Text.assemble(header, "\n\n", content), + title=f"[bold blue]Tool Output (ID: {call_id})[/bold blue]", + subtitle="[bold green]Live Update[/bold green]", + border_style="blue", + padding=(1, 2), + box=ROUNDED + ) + + # 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 blue]Tool Output: {tool_call}[/bold blue]", + subtitle=subtitle, + border_style="blue", padding=(1, 2), - title="[bold]Tool Output[/bold]", - title_align="left", - expand=True + box=ROUNDED ) - console.print(output_panel) \ No newline at end of file + + # 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) + # This works better with streamed content as it doesn't scroll the terminal + 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 From 53f433f7d3eecdae23ee26a712a533f25095f3cb Mon Sep 17 00:00:00 2001 From: lidia9 Date: Fri, 25 Apr 2025 12:35:56 +0200 Subject: [PATCH 8/9] tool output streaming works --- .../agents/models/openai_chatcompletions.py | 81 ++++++++++--------- src/cai/tools/common.py | 38 ++++++--- .../reconnaissance/generic_linux_command.py | 13 ++- src/cai/util.py | 28 ++++--- 4 files changed, 103 insertions(+), 57 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index f658e183..e294bee7 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -1609,43 +1609,50 @@ class _Converter: # Display the tool output immediately with the matched tool call from cai.util import cli_print_tool_output - # 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, - execution_info=execution_info, - token_info=token_info - ) + # 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() diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index ebf6821e..91d9dbbe 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -245,6 +245,14 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None): 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) @@ -288,10 +296,16 @@ def _run_local_streamed(command, call_id, timeout=100): # 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") - header.append(f"command='{command}'", 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() @@ -299,9 +313,9 @@ def _run_local_streamed(command, call_id, timeout=100): panel = Panel( Text.assemble(header, "\n\n", content), - title="[bold blue]Tool Execution[/bold blue]", + title="[bold green]Tool Execution[/bold green]", subtitle="[bold green]Live Output[/bold green]", - border_style="blue", + border_style="green", padding=(1, 2), box=ROUNDED ) @@ -320,9 +334,9 @@ def _run_local_streamed(command, call_id, timeout=100): content.append(line, style="bright_white") panel = Panel( Text.assemble(header, "\n\n", content), - title="[bold blue]Tool Execution[/bold blue]", + title="[bold green]Tool Execution[/bold green]", subtitle="[bold green]Live Output[/bold green]", - border_style="blue", + border_style="green", padding=(1, 2), box=ROUNDED ) @@ -340,9 +354,9 @@ def _run_local_streamed(command, call_id, timeout=100): output_buffer.append("\nERROR OUTPUT:\n" + stderr_data) panel = Panel( Text.assemble(header, "\n\n", content), - title="[bold blue]Tool Execution[/bold blue]", + title="[bold green]Tool Execution[/bold green]", subtitle="[bold green]Live Output[/bold green]", - border_style="blue", + border_style="green", padding=(1, 2), box=ROUNDED ) @@ -353,9 +367,9 @@ def _run_local_streamed(command, call_id, timeout=100): content.append(f"\nCommand {completion_status}", style="green") panel = Panel( Text.assemble(header, "\n\n", content), - title="[bold blue]Tool Execution[/bold blue]", + title="[bold green]Tool Execution[/bold green]", subtitle=f"[bold green]{completion_status}[/bold green]", - border_style="blue", + border_style="green", padding=(1, 2), box=ROUNDED ) @@ -365,7 +379,11 @@ def _run_local_streamed(command, call_id, timeout=100): time.sleep(0.5) else: # Fallback to simpler streaming with cli_print_tool_output - tool_args = {"command": command} + # 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) diff --git a/src/cai/tools/reconnaissance/generic_linux_command.py b/src/cai/tools/reconnaissance/generic_linux_command.py index 7fa156e7..6c1274b8 100644 --- a/src/cai/tools/reconnaissance/generic_linux_command.py +++ b/src/cai/tools/reconnaissance/generic_linux_command.py @@ -90,9 +90,20 @@ def generic_linux_command(command: str = "", 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] # Run the command with the appropriate parameters - return run_command(full_command, ctf=ctf, + result = run_command(full_command, ctf=ctf, async_mode=async_mode, session_id=session_id, 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 6397a230..087c98d2 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -877,8 +877,11 @@ def cli_print_agent_messages(agent_name, message, counter, model, debug, # pyli console.print(panel) # If there are tool panels, print them after the main message panel - for tool_panel in tool_panels: - console.print(tool_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.""" @@ -1072,10 +1075,17 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut 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: @@ -1127,11 +1137,12 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut # Create the panel for display panel = Panel( Text.assemble(header, "\n\n", content), - title=f"[bold blue]Tool Output (ID: {call_id})[/bold blue]", - subtitle="[bold green]Live Update[/bold green]", + 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 + box=ROUNDED, + title_align="left" ) # Display using Rich @@ -1181,9 +1192,9 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut # Create the panel with all sections panel = Panel( Text.assemble(*sections), - title=f"[bold blue]Tool Output: {tool_call}[/bold blue]", + title=f"[bold green]Tool Output: {tool_call}[/bold green]", subtitle=subtitle, - border_style="blue", + border_style="green", padding=(1, 2), box=ROUNDED ) @@ -1215,7 +1226,6 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut header = header[:term_width-4] # Clear the screen for the tool output (alternative approach) - # This works better with streamed content as it doesn't scroll the terminal print(f"\r{header}") # Print the content From 791fbc5e967f534f158bc9b08435c48337aaee81 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Fri, 25 Apr 2025 12:41:54 +0200 Subject: [PATCH 9/9] fix ctrl+c --- src/cai/cli.py | 20 +++++++++++--------- 1 file changed, 11 insertions(+), 9 deletions(-) diff --git a/src/cai/cli.py b/src/cai/cli.py index 1d2fed05..72940444 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -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() @@ -358,14 +360,14 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= #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: