From 58872b66679fc350e713ec3f83fe1d732f1eb866 Mon Sep 17 00:00:00 2001 From: lidia9 Date: Tue, 22 Apr 2025 15:12:49 +0000 Subject: [PATCH] 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