From 56e6fcf1a411a6e48fcc3aaf0eca0a6e17f812bb Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Thu, 24 Apr 2025 15:03:13 +0200 Subject: [PATCH 01/29] fix agent command --- examples/cai/.gitkeep | 0 src/cai/cli.py | 5 +- src/cai/repl/commands/agent.py | 201 +++++++++++++++------------------ src/cai/util.py | 91 +++++++-------- 4 files changed, 130 insertions(+), 167 deletions(-) delete mode 100644 examples/cai/.gitkeep diff --git a/examples/cai/.gitkeep b/examples/cai/.gitkeep deleted file mode 100644 index e69de29b..00000000 diff --git a/src/cai/cli.py b/src/cai/cli.py index 53c9ce05..bdbb54e7 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -280,8 +280,9 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= if commands_handle_command(command, args): continue # Command was handled, continue to next iteration - # If command wasn't recognized, show error - console.print(f"[red]Unknown command: {command}[/red]") + # If command wasn't recognized, show error (skip for /shell or /s) + if command not in ("/shell", "/s"): + console.print(f"[red]Unknown command: {command}[/red]") continue # Process the conversation with the agent diff --git a/src/cai/repl/commands/agent.py b/src/cai/repl/commands/agent.py index 046d4d3e..b23194a3 100644 --- a/src/cai/repl/commands/agent.py +++ b/src/cai/repl/commands/agent.py @@ -155,50 +155,45 @@ class AgentCommand(Command): table = Table(title="Available Agents") table.add_column("#", style="dim") table.add_column("Name", style="cyan") - table.add_column("Module", style="magenta") + table.add_column("Key", style="magenta") + table.add_column("Module", style="green") table.add_column("Description", style="green") - table.add_column("Pattern", style="blue") - table.add_column("Model", style="yellow") - # Scan all agents from the agents folder + # Retrieve all registered agents agents_to_display = get_available_agents() - # Display all agents - for i, (name, agent) in enumerate(agents_to_display.items(), 1): - description = agent.description - if not description and hasattr(agent, 'instructions'): - if callable(agent.instructions): - description = agent.instructions(context_variables={}) - else: - description = agent.instructions - # Clean up description - remove newlines and strip spaces + for idx, (agent_key, agent) in enumerate(agents_to_display.items(), start=1): + # Human-friendly name (falls back to the dict key) + display_name = getattr(agent, "name", agent_key) + + # Use provided description, otherwise derive from instructions + description = getattr(agent, "description", "") or "" + if not description and hasattr(agent, "instructions"): + instr = agent.instructions + description = instr(context_variables={}) if callable(instr) else instr if isinstance(description, str): description = " ".join(description.split()) if len(description) > 50: description = description[:47] + "..." - # Get the module name for the agent - module_name = get_agent_module(name) + # Module where this agent lives + module_name = get_agent_module(agent_key) - # Get the pattern if it exists - pattern = getattr(agent, 'pattern', '') - if pattern: - pattern = pattern.capitalize() - - # Handle model display based on agent type - model_display = self._get_model_display(name, agent) + # Add a row with all collected info table.add_row( - str(i), - name, + str(idx), + display_name, + agent_key, module_name, - description, - pattern, - model_display + description ) console.print(table) return True + console.print(table) + return True + def handle_select(self, args: Optional[List[str]] = None) -> bool: # pylint: disable=too-many-branches,line-too-long # noqa: E501 """Handle /agent select command. @@ -210,76 +205,41 @@ class AgentCommand(Command): """ if not args: console.print("[red]Error: No agent specified[/red]") - console.print("Usage: /agent select ") + console.print("Usage: /agent select ") return False agent_id = args[0] - # Get the list of available agents agents_to_display = get_available_agents() + agent_list = list(agents_to_display.items()) # preserve order for indexing # Check if agent_id is a number if agent_id.isdigit(): index = int(agent_id) - if 1 <= index <= len(agents_to_display): - agent_name = list(agents_to_display.keys())[index - 1] + if 1 <= index <= len(agent_list): + # Get the agent tuple from the list + selected_agent_key, selected_agent = agent_list[index - 1] + agent_name = getattr(selected_agent, "name", selected_agent_key) + agent = selected_agent else: - console.print( - f"[red]Error: Invalid agent number: {agent_id}[/red]") + console.print(f"[red]Error: Invalid agent number: {agent_id}[/red]") return False else: - # Treat as agent name - agent_name = agent_id - if agent_name not in agents_to_display: - console.print(f"[red]Error: Unknown agent: {agent_name}[/red]") + # Treat as agent key + for key, agent_obj in agents_to_display.items(): + if key == agent_id: + agent = agent_obj + agent_name = getattr(agent_obj, "name", key) + break + else: + console.print(f"[red]Error: Unknown agent key: {agent_id}[/red]") return False - # Get the agent - agent = agents_to_display[agent_name] - - # Set the agent as the current agent in the REPL - # We need to avoid circular imports, so we'll use a different approach - # to access the client and current_agent variables - - # Import the module dynamically to avoid circular imports - if 'cai.repl.repl' in sys.modules: - repl_module = sys.modules['cai.repl.repl'] - - # Check if client is initialized - if hasattr(repl_module, 'client') and repl_module.client: - # Update the active_agent in the client - repl_module.client.active_agent = agent - - # Update the global current_agent variable if it exists - if hasattr(repl_module, 'current_agent'): - repl_module.current_agent = agent - - # Update the global agent variable if it exists - if hasattr(repl_module, 'agent'): - repl_module.agent = agent - - # Also update the agent variable in the run_demo_loop - # function's frame if possible - try: - for frame_info in inspect.stack(): - frame = frame_info.frame - if ('run_demo_loop' in frame.f_code.co_name and - 'agent' in frame.f_locals): - frame.f_locals['agent'] = agent - break - except Exception: # pylint: disable=broad-except # nosec - # If this fails, we still have the global current_agent as - # a fallback - pass - - console.print( - f"[green]Switched to agent: {agent_name}[/green]") - visualize_agent_graph(agent) - return True - console.print("[red]Error: CAI client not initialized[/red]") - return False - console.print("[red]Error: REPL module not initialized[/red]") - return False + os.environ["CAI_MODEL"] = agent_name + console.print( + f"[green]Switched to agent: {agent_name}[/green]") + visualize_agent_graph(agent) + return True def handle_info(self, args: Optional[List[str]] = None) -> bool: """Handle /agent info command. @@ -292,57 +252,74 @@ class AgentCommand(Command): """ if not args: console.print("[red]Error: No agent specified[/red]") - console.print("Usage: /agent info ") + console.print("Usage: /agent info ") return False agent_id = args[0] - # Get the list of available agents + # Get available agents agents_to_display = get_available_agents() - # Check if agent_id is a number + # Resolve agent_id to an agent key (by index or name) if agent_id.isdigit(): - index = int(agent_id) - if 1 <= index <= len(agents_to_display): - agent_name = list(agents_to_display.keys())[index - 1] - else: - console.print( - f"[red]Error: Invalid agent number: {agent_id}[/red]") + idx = int(agent_id) + if not (1 <= idx <= len(agents_to_display)): + console.print(f"[red]Error: Invalid agent number: {agent_id}[/red]") return False + agent_key = list(agents_to_display.keys())[idx - 1] else: - # Treat as agent name - agent_name = agent_id - if agent_name not in agents_to_display: - console.print(f"[red]Error: Unknown agent: {agent_name}[/red]") + agent_key = None + for key, ag in agents_to_display.items(): + if key == agent_id or getattr(ag, "name", "").lower() == agent_id.lower(): + agent_key = key + break + if agent_key is None: + console.print(f"[red]Error: Unknown agent key: {agent_id}[/red]") return False - # Get the agent - agent = agents_to_display[agent_name] + agent = agents_to_display[agent_key] - # Display agent information + # Display agent information instructions = agent.instructions if callable(instructions): instructions = instructions() + # Prepare agent properties + name = agent.name or agent_key + description = getattr(agent, "description", None) or "N/A" + clean_description = " ".join(line.strip() for line in description.splitlines()) + functions = getattr(agent, "functions", []) + parallel = getattr(agent, "parallel_tool_calls", False) + handoff_desc = getattr(agent, "handoff_description", None) or "N/A" + handoffs = getattr(agent, "handoffs", []) + tools = getattr(agent, "tools", []) + guardrails_in = getattr(agent, "input_guardrails", []) + guardrails_out = getattr(agent, "output_guardrails", []) + output_type = getattr(agent, "output_type", None) or "N/A" + hooks = getattr(agent, "hooks", []) or [] - # Handle model display based on agent type - model_display = self._get_model_display_for_info(agent_name, agent) - - # Create a markdown table for agent details + # Build markdown content for agent info markdown_content = f""" -# Agent: {agent_name} +# Agent Info: {name} -| Property | Value | -|----------|-------| -| Name | {agent.name} | -| Model | {model_display} | -| Functions | {len(agent.functions)} | -| Parallel Tool Calls | {'Yes' if agent.parallel_tool_calls else 'No'} | +| Property | Value | +|------------------------|-------------------------------| +| Key | {agent_key} | +| Name | {name} | +| Description | {clean_description} | +| Functions | {len(functions)} | +| Parallel Tool Calls | {"Yes" if parallel else "No"} | +| Handoff Description | {handoff_desc} | +| Handoffs | {len(handoffs)} | +| Tools | {len(tools)} | +| Input Guardrails | {len(guardrails_in)} | +| Output Guardrails | {len(guardrails_out)} | +| Output Type | {output_type} | +| Hooks | {len(hooks)} | ## Instructions - {instructions} -""" +""" console.print(Markdown(markdown_content)) return True diff --git a/src/cai/util.py b/src/cai/util.py index 63cedaf5..81703c68 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -268,87 +268,72 @@ def load_prompt_template(template_path): except Exception as e: raise ValueError(f"Failed to load template '{template_path}': {str(e)}") +# Start of Selection def visualize_agent_graph(start_agent): """ Visualize agent graph showing all bidirectional connections between agents. Uses Rich library for pretty printing. """ - console = Console() # pylint: disable=redefined-outer-name + console = Console() if start_agent is None: console.print("[red]No agent provided to visualize.[/red]") return - tree = Tree( - f"🤖 { - start_agent.name} (Current Agent)", - guide_style="bold blue") + tree = Tree(f"🤖 {start_agent.name} (Current Agent)", guide_style="bold blue") - # Track visited agents and their nodes to handle cross-connections - visited = {} + visited = set() agent_nodes = {} - agent_positions = {} # Track positions in tree - position_counter = 0 # Counter for tracking positions + agent_positions = {} + position_counter = 0 - def add_agent_node(agent, parent=None, is_transfer=False): # pylint: disable=too-many-branches # noqa: E501 - """Add agent node and track for cross-connections""" + def add_agent_node(agent, parent=None, is_transfer=False): + """Add an agent node and track for cross-connections.""" nonlocal position_counter - if agent is None: return None + aid = id(agent) + if aid in visited: + if is_transfer and parent: + original_pos = agent_positions.get(aid) + parent.add(f"[cyan]↩ Return to {agent.name} (Agent #{original_pos})[/cyan]") + return agent_nodes.get(aid) - # Create or get existing node for this agent - if id(agent) in visited: - if is_transfer: - # Add reference with position for repeated agents - original_pos = agent_positions[id(agent)] - parent.add( - f"[cyan]↩ Return to { - agent.name} (Top Level Agent #{original_pos})[/cyan]") - return agent_nodes[id(agent)] - - visited[id(agent)] = True + visited.add(aid) position_counter += 1 - agent_positions[id(agent)] = position_counter + agent_positions[aid] = position_counter - # Create node for current agent - if is_transfer: + if is_transfer and parent: node = parent + elif parent: + node = parent.add(f"[green]{agent.name} (#{position_counter})[/green]") else: - node = parent.add( - f"[green]{agent.name} (#{position_counter})[/green]") if parent else tree # noqa: E501 pylint: disable=line-too-long - agent_nodes[id(agent)] = node + node = tree + agent_nodes[aid] = node - # Add tools as children + # Add tools tools_node = node.add("[yellow]Tools[/yellow]") - for fn in getattr(agent, "functions", []): - if callable(fn): - fn_name = getattr(fn, "__name__", "") - if ("handoff" not in fn_name.lower() and - not fn_name.startswith("transfer_to")): - tools_node.add(f"[blue]{fn_name}[/blue]") + for tool in getattr(agent, "tools", []): + tool_name = getattr(tool, "name", None) or getattr(tool, "__name__", "") + tools_node.add(f"[blue]{tool_name}[/blue]") - # Add Handoffs section + # Add handoffs transfers_node = node.add("[magenta]Handoffs[/magenta]") + for handoff_fn in getattr(agent, "handoffs", []): + if callable(handoff_fn): + try: + next_agent = handoff_fn() + if next_agent: + transfer_node = transfers_node.add(f"🤖 {next_agent.name}") + add_agent_node(next_agent, transfer_node, True) + except Exception: + continue - # Process handoff functions - for fn in getattr(agent, "functions", []): # pylint: disable=too-many-nested-blocks # noqa: E501 - if callable(fn): - fn_name = getattr(fn, "__name__", "") - if ("handoff" in fn_name.lower() or - fn_name.startswith("transfer_to")): - try: - next_agent = fn() - if next_agent: - # Show bidirectional connection - transfer = transfers_node.add( - f"🤖 {next_agent.name}") # noqa: E501 - add_agent_node(next_agent, transfer, True) - except Exception: # nosec: B112 # pylint: disable=broad-exception-caught # noqa: E501 - continue return node - # Start recursive traversal from root agent + + # Start traversal from the root agent add_agent_node(start_agent) console.print(tree) +# End of Selectio def fix_litellm_transcription_annotations(): """ From 093d6710deb99c67783cf9550c80f793b5acb39e Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Thu, 24 Apr 2025 15:05:47 +0200 Subject: [PATCH 02/29] fix agent command 2 --- src/cai/repl/commands/agent.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cai/repl/commands/agent.py b/src/cai/repl/commands/agent.py index b23194a3..e01d4fdf 100644 --- a/src/cai/repl/commands/agent.py +++ b/src/cai/repl/commands/agent.py @@ -211,7 +211,7 @@ class AgentCommand(Command): agent_id = args[0] agents_to_display = get_available_agents() - agent_list = list(agents_to_display.items()) # preserve order for indexing + agent_list = list(agents_to_display.items()) # Check if agent_id is a number if agent_id.isdigit(): From e2e54b2008216111d1185447d1abb87ca2d2c34a Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Fri, 25 Apr 2025 10:24:58 +0200 Subject: [PATCH 03/29] fix history command --- src/cai/repl/commands/history.py | 6 +- .../agents/models/openai_chatcompletions.py | 108 +++++++++++++++++- 2 files changed, 107 insertions(+), 7 deletions(-) diff --git a/src/cai/repl/commands/history.py b/src/cai/repl/commands/history.py index 3a7df9c9..d776d766 100644 --- a/src/cai/repl/commands/history.py +++ b/src/cai/repl/commands/history.py @@ -29,13 +29,13 @@ class HistoryCommand(Command): """ # Access messages directly from repl.py's global scope try: - from cai.repl.repl import messages # pylint: disable=import-outside-toplevel # noqa: E501 + from cai.sdk.agents.models.openai_chatcompletions import message_history # pylint: disable=import-outside-toplevel # noqa: E501 except ImportError: console.print( "[red]Error: Could not access conversation history[/red]") return False - if not messages: + if not message_history: console.print("[yellow]No conversation history available[/yellow]") return True @@ -50,7 +50,7 @@ class HistoryCommand(Command): table.add_column("Content", style="green") # Add messages to the table - for idx, msg in enumerate(messages, 1): + for idx, msg in enumerate(message_history, 1): role = msg.get("role", "unknown") content = msg.get("content", "") diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 08d5dfe2..53b61eb6 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -96,6 +96,7 @@ litellm.suppress_debug_info = True _USER_AGENT = f"Agents/Python {__version__}" _HEADERS = {"User-Agent": _USER_AGENT} +message_history = [] @dataclass class _StreamingState: @@ -235,7 +236,28 @@ class OpenAIChatCompletionsModel(Model): "role": "system", }, ) - + # --- Add to message_history: user, system, and assistant tool call messages --- + # Add system prompt to message_history + if system_instructions: + message_history.append({ + "role": "system", + "content": system_instructions + }) + # Add user prompt(s) to message_history + if isinstance(input, str): + message_history.append({ + "role": "user", + "content": input + }) + elif isinstance(input, list): + for item in input: + # Try to extract user messages + if isinstance(item, dict): + if item.get("role") == "user": + message_history.append({ + "role": "user", + "content": item.get("content", "") + }) # Get token count estimate before API call for consistent counting estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages) @@ -335,6 +357,34 @@ class OpenAIChatCompletionsModel(Model): tool_output=None, # Don't pass tool output here, we're using direct display ) + # --- Add assistant tool call to message_history if present --- + # If the response contains tool_calls, add them to message_history as assistant messages + assistant_msg = response.choices[0].message + if hasattr(assistant_msg, "tool_calls") and assistant_msg.tool_calls: + for tool_call in assistant_msg.tool_calls: + # Compose a message for the tool call + tool_call_msg = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": tool_call.id, + "type": tool_call.type, + "function": { + "name": tool_call.function.name, + "arguments": tool_call.function.arguments + } + } + ] + } + message_history.append(tool_call_msg) + # If the assistant message is just text, add it as well + elif hasattr(assistant_msg, "content") and assistant_msg.content: + message_history.append({ + "role": "assistant", + "content": assistant_msg.content + }) + usage = ( Usage( requests=1, @@ -404,7 +454,25 @@ class OpenAIChatCompletionsModel(Model): "role": "system", }, ) - + # --- Add to message_history: user, system prompts --- + if system_instructions: + message_history.append({ + "role": "system", + "content": system_instructions + }) + if isinstance(input, str): + message_history.append({ + "role": "user", + "content": input + }) + elif isinstance(input, list): + for item in input: + if isinstance(item, dict): + if item.get("role") == "user": + message_history.append({ + "role": "user", + "content": item.get("content", "") + }) # Get token count estimate before API call for consistent counting estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages) @@ -429,7 +497,9 @@ class OpenAIChatCompletionsModel(Model): # Initialize a streaming text accumulator for rich display streaming_text_buffer = "" - + # For tool call streaming, accumulate tool_calls to add to message_history at the end + streamed_tool_calls = [] + async for chunk in stream: if not state.started: state.started = True @@ -639,6 +709,26 @@ class OpenAIChatCompletionsModel(Model): state.function_calls[tc_index].call_id += call_id + # --- Accumulate tool call for message_history --- + # Only add if not already present (avoid duplicates in streaming) + tool_call_msg = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": state.function_calls[tc_index].call_id, + "type": "function", + "function": { + "name": state.function_calls[tc_index].name, + "arguments": state.function_calls[tc_index].arguments + } + } + ] + } + # Only add if not already in streamed_tool_calls + if tool_call_msg not in streamed_tool_calls: + streamed_tool_calls.append(tool_call_msg) + function_call_starting_index = 0 if state.text_content_index_and_output: function_call_starting_index += 1 @@ -838,7 +928,17 @@ class OpenAIChatCompletionsModel(Model): total_cost=total_cost, ) break - + + # --- Add assistant tool call(s) to message_history at the end of streaming --- + for tool_call_msg in streamed_tool_calls: + message_history.append(tool_call_msg) + # If there was only text output, add that as an assistant message + if (not streamed_tool_calls) and state.text_content_index_and_output and state.text_content_index_and_output[1].text: + message_history.append({ + "role": "assistant", + "content": state.text_content_index_and_output[1].text + }) + if tracing.include_data(): span_generation.span_data.output = [final_response.model_dump()] From e561a5b59ad3656db4d46ba9aa0801369b69c69e Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Fri, 25 Apr 2025 14:07:32 +0200 Subject: [PATCH 04/29] fix config for CAI_STREAM --- src/cai/repl/commands/config.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/cai/repl/commands/config.py b/src/cai/repl/commands/config.py index 81184e6d..3588cd19 100644 --- a/src/cai/repl/commands/config.py +++ b/src/cai/repl/commands/config.py @@ -122,6 +122,11 @@ ENV_VARS = { "name": "CAI_SUPPORT_INTERVAL", "description": "Number of turns between support agent executions", "default": "5" + }, + 22: { + "name": "CAI_STREAM", + "description": "Boolean to enable real-time, chunked responses instead of full messages.", + "default": "True" }, } From 6973e23ece9319c9dbd54db25b0a02041e0ed6f7 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Fri, 25 Apr 2025 14:12:57 +0200 Subject: [PATCH 05/29] fix history error --- src/cai/repl/commands/history.py | 44 +++++++++++++++++--------------- 1 file changed, 24 insertions(+), 20 deletions(-) diff --git a/src/cai/repl/commands/history.py b/src/cai/repl/commands/history.py index d776d766..38eaeb96 100644 --- a/src/cai/repl/commands/history.py +++ b/src/cai/repl/commands/history.py @@ -51,30 +51,34 @@ class HistoryCommand(Command): # Add messages to the table for idx, msg in enumerate(message_history, 1): - role = msg.get("role", "unknown") - content = msg.get("content", "") + try: + role = msg.get("role", "unknown") + content = msg.get("content", "") - # Truncate long content for better display - if len(content) > 100: - content = content[:97] + "..." + # Truncate long content for better display + if len(content) > 100: + content = content[:97] + "..." - # Color the role based on type - if role == "user": - role_style = "cyan" - elif role == "assistant": - role_style = "yellow" - else: - role_style = "red" + # Color the role based on type + if role == "user": + role_style = "cyan" + elif role == "assistant": + role_style = "yellow" + else: + role_style = "red" - # Add a newline between each role for better readability - if idx > 1: - table.add_row("", "", "") + # Add a newline between each role for better readability + if idx > 1: + table.add_row("", "", "") - table.add_row( - str(idx), - f"[{role_style}]{role}[/{role_style}]", - content - ) + table.add_row( + str(idx), + f"[{role_style}]{role}[/{role_style}]", + content + ) + except Exception: + # Si hay un error, evita esto (no agregues la fila) + continue console.print(table) return True From 993ab6f27e374923d396c0cb22edb7be812ad110 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Fri, 25 Apr 2025 14:30:37 +0200 Subject: [PATCH 06/29] ad gpt4o-mini --- src/cai/repl/commands/model.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/cai/repl/commands/model.py b/src/cai/repl/commands/model.py index 8595ede7..04d2ad52 100644 --- a/src/cai/repl/commands/model.py +++ b/src/cai/repl/commands/model.py @@ -145,6 +145,9 @@ class ModelCommand(Command): {"name": "gpt-4-turbo", "description": "Fast and powerful GPT-4 model"} ], + "OpenAI GPT-4o-mini": [ + {"name": "gpt-4o-mini", "description": " GPT-4o mini model"} + ], "OpenAI GPT-4.5": [ { "name": "gpt-4.5-preview", From ee88d5e78423a371d413f7c8c2c64a69e88f5067 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 5 May 2025 11:10:19 +0200 Subject: [PATCH 07/29] fix agent streaming --- src/cai/util.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/cai/util.py b/src/cai/util.py index 81703c68..0eba2afe 100644 --- a/src/cai/util.py +++ b/src/cai/util.py @@ -990,6 +990,7 @@ def update_agent_streaming_content(context, text_delta): # Force an update with the new panel context["live"].update(updated_panel) context["panel"] = updated_panel + context["live"].refresh() def finish_agent_streaming(context, final_stats=None): """Finish the streaming session and display final stats if available.""" From 057dbc2aa0785f5964033c59ccbfc533c32b8852 Mon Sep 17 00:00:00 2001 From: lidia9 Date: Mon, 5 May 2025 13:19:42 +0200 Subject: [PATCH 08/29] FIX for making o3-mini work --- src/cai/sdk/agents/models/openai_chatcompletions.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 53b61eb6..8ef8df3a 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -93,6 +93,9 @@ if TYPE_CHECKING: # Suppress debug info from litellm litellm.suppress_debug_info = True +if os.getenv('CAI_MODEL') == "o3-mini": + litellm.drop_params = True + _USER_AGENT = f"Agents/Python {__version__}" _HEADERS = {"User-Agent": _USER_AGENT} From 716850a643420656f1e7d504fd762035f60ba02b Mon Sep 17 00:00:00 2001 From: lidia9 Date: Mon, 5 May 2025 13:31:15 +0200 Subject: [PATCH 09/29] FIX litellm.drop_params=true for gemini --- src/cai/sdk/agents/models/openai_chatcompletions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 8ef8df3a..9b4579d8 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -93,7 +93,7 @@ if TYPE_CHECKING: # Suppress debug info from litellm litellm.suppress_debug_info = True -if os.getenv('CAI_MODEL') == "o3-mini": +if os.getenv('CAI_MODEL') == "o3-mini" or os.getenv('CAI_MODEL') == "gemini-1.5-pro": litellm.drop_params = True _USER_AGENT = f"Agents/Python {__version__}" From fe20d4bfef9ab638b7e4d7ae592f02aff7f9be38 Mon Sep 17 00:00:00 2001 From: lidia9 Date: Mon, 5 May 2025 15:21:39 +0200 Subject: [PATCH 10/29] tool_call execution and parsing for qwen --- .../agents/models/openai_chatcompletions.py | 318 ++++++++++++++---- 1 file changed, 255 insertions(+), 63 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 9b4579d8..f1ee8921 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -7,6 +7,7 @@ import os import litellm import tiktoken import inspect +import hashlib from collections.abc import AsyncIterator, Iterable from dataclasses import dataclass, field @@ -526,6 +527,10 @@ class OpenAIChatCompletionsModel(Model): choices = [{"delta": chunk.delta}] elif isinstance(chunk, dict) and 'choices' in chunk: choices = chunk['choices'] + # Special handling for Qwen/Ollama chunks + elif isinstance(chunk, dict) and ('content' in chunk or 'function_call' in chunk): + # Qwen direct delta format - convert to standard + choices = [{"delta": chunk}] else: # Skip chunks that don't contain choice data continue @@ -662,11 +667,7 @@ class OpenAIChatCompletionsModel(Model): # Handle tool calls # Because we don't know the name of the function until the end of the stream, we'll # save everything and yield events at the end - tool_calls = None - if hasattr(delta, 'tool_calls') and delta.tool_calls: - tool_calls = delta.tool_calls - elif isinstance(delta, dict) and 'tool_calls' in delta and delta['tool_calls']: - tool_calls = delta['tool_calls'] + tool_calls = self._detect_and_format_function_calls(delta) if tool_calls: for tc_delta in tool_calls: @@ -709,7 +710,12 @@ class OpenAIChatCompletionsModel(Model): call_id = tc_delta.id or "" elif isinstance(tc_delta, dict) and 'id' in tc_delta: call_id = tc_delta.get('id', "") or "" - + else: + # For Qwen models, generate a predictable ID if none is provided + if state.function_calls[tc_index].name: + # Generate a stable ID from the function name and arguments + call_id = f"call_{hashlib.md5(state.function_calls[tc_index].name.encode()).hexdigest()[:8]}" + state.function_calls[tc_index].call_id += call_id # --- Accumulate tool call for message_history --- @@ -1067,31 +1073,69 @@ class OpenAIChatCompletionsModel(Model): "extra_headers": _HEADERS, } - # Model adjustments - if any(x in self.model for x in ["claude"]): - litellm.drop_params = True - - # Error encountered: Error code: 400 - {'error': {'code': 'invalid_request_error', - # 'message': "'tool_choice' is only allowed when 'tools' are specified", - # 'type': 'invalid_request_error', 'param': None}} - # - # if has no tools, remove tool_choice - if not converted_tools: - kwargs.pop("tool_choice", None) - - # BadRequestError encountered: litellm.BadRequestError: AnthropicException - - # b'{"type":"error","error": - # {"type":"invalid_request_error","message":"store: Extra inputs are not permitted"}}' - # - kwargs.pop("store", None) + # Determine provider based on model string + model_str = str(self.model).lower() + + # Provider-specific adjustments + if "/" in model_str: + # Handle provider/model format + provider = model_str.split("/")[0] - # Filter out NotGiven values to avoid JSON serialization issues - filtered_kwargs = {} - for key, value in kwargs.items(): - if value is not NOT_GIVEN: - filtered_kwargs[key] = value - kwargs = filtered_kwargs + # Apply provider-specific configurations + if provider == "deepseek": + litellm.drop_params = True + kwargs.pop("parallel_tool_calls", None) + # Remove tool_choice if no tools are specified + if not converted_tools: + kwargs.pop("tool_choice", None) + elif provider == "claude": + litellm.drop_params = True + kwargs.pop("store", None) + # Remove tool_choice if no tools are specified + if not converted_tools: + kwargs.pop("tool_choice", None) + elif provider == "gemini": + kwargs.pop("parallel_tool_calls", None) + # Add any specific gemini settings if needed + else: + # Handle models without provider prefix + if "claude" in model_str: + litellm.drop_params = True + # Remove store parameter which isn't supported by Anthropic + kwargs.pop("store", None) + # Remove tool_choice if no tools are specified + if not converted_tools: + kwargs.pop("tool_choice", None) + elif "gemini" in model_str: + kwargs.pop("parallel_tool_calls", None) + elif "qwen" in model_str or ":" in model_str: + # Handle Ollama-served models with custom formats (e.g., qwen2.5:14b) + # These typically need the Ollama provider + litellm.drop_params = True + kwargs.pop("parallel_tool_calls", None) + # These models may not support certain parameters + if not converted_tools: + kwargs.pop("tool_choice", None) + # Don't add custom_llm_provider here to avoid duplication with Ollama provider + if self.is_ollama: + # Clean kwargs for ollama to avoid parameter conflicts + for param in ["custom_llm_provider"]: + kwargs.pop(param, None) + elif any(x in model_str for x in ["o1", "o3", "o4"]): + # Handle OpenAI reasoning models (o1, o3, o4) + kwargs.pop("parallel_tool_calls", None) + # Add reasoning effort if provided + if hasattr(model_settings, "reasoning_effort"): + kwargs["reasoning_effort"] = model_settings.reasoning_effort + + # Filter out NotGiven values to avoid JSON serialization issues + filtered_kwargs = {} + for key, value in kwargs.items(): + if value is not NOT_GIVEN: + filtered_kwargs[key] = value + kwargs = filtered_kwargs + try: if self.is_ollama: return await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) @@ -1101,21 +1145,70 @@ class OpenAIChatCompletionsModel(Model): except litellm.exceptions.BadRequestError as e: # print(color("BadRequestError encountered: " + str(e), fg="yellow")) if "LLM Provider NOT provided" in str(e): - # Create a copy of params to avoid overwriting the original - # ones - try: - return await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) - except litellm.exceptions.BadRequestError as e: # pylint: disable=W0621,C0301 # noqa: E501 - # - # CTRL-C handler for ollama models - # - if "invalid message content type" in str(e): - kwargs["messages"] = fix_message_list( - kwargs["messages"]) + model_str = str(self.model).lower() + provider = None + is_qwen = "qwen" in model_str or ":" in model_str + + # Special handling for Qwen models + if is_qwen: + try: + # Use the specialized Qwen approach first return await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) + except Exception as qwen_e: + # If that fails, try our direct OpenAI approach + qwen_params = kwargs.copy() + qwen_params["api_base"] = get_ollama_api_base() + qwen_params["custom_llm_provider"] = "openai" # Use openai provider + + # Make sure tools are passed + if "tools" in kwargs and kwargs["tools"]: + qwen_params["tools"] = kwargs["tools"] + if "tool_choice" in kwargs and kwargs["tool_choice"] is not NOT_GIVEN: + qwen_params["tool_choice"] = kwargs["tool_choice"] + + try: + if stream: + # Streaming case + response = Response( + id=FAKE_RESPONSES_ID, + created_at=time.time(), + model=self.model, + object="response", + output=[], + tool_choice="auto" if tool_choice is None or tool_choice == NOT_GIVEN else cast(Literal["auto", "required", "none"], tool_choice), + top_p=model_settings.top_p, + temperature=model_settings.temperature, + tools=[], + parallel_tool_calls=parallel_tool_calls or False, + ) + stream_obj = await litellm.acompletion(**qwen_params) + return response, stream_obj + else: + # Non-streaming case + ret = litellm.completion(**qwen_params) + return ret + except Exception as direct_e: + # All approaches failed, log and raise the original error + print(f"All Qwen approaches failed. Original error: {str(e)}, Direct error: {str(direct_e)}") + raise e + + # Try to detect provider from model string + if "/" in model_str: + provider = model_str.split("/")[0] + + if provider: + # Add provider-specific settings based on detected provider + provider_kwargs = kwargs.copy() + if provider == "deepseek": + provider_kwargs["custom_llm_provider"] = "deepseek" + elif provider == "claude" or "claude" in model_str: + provider_kwargs["custom_llm_provider"] = "anthropic" + elif provider == "gemini": + provider_kwargs["custom_llm_provider"] = "gemini" else: - raise e - + # For unknown providers, try ollama as fallback + return await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) + elif ("An assistant message with 'tool_calls'" in str(e) or "`tool_use` blocks must be followed by a user message with `tool_result`" in str(e)): # noqa: E501 # pylint: disable=C0301 print(f"Error: {str(e)}") @@ -1228,25 +1321,48 @@ class OpenAIChatCompletionsModel(Model): stream: bool, parallel_tool_calls: bool ) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]: - # Filter out parameters not supported by Ollama ollama_supported_params = { - "model": kwargs["model"], - "messages": kwargs["messages"], - "temperature": kwargs["temperature"] if kwargs["temperature"] is not NOT_GIVEN else None, - "top_p": kwargs["top_p"] if kwargs["top_p"] is not NOT_GIVEN else None, - "max_tokens": kwargs["max_tokens"] if kwargs["max_tokens"] is not NOT_GIVEN else None, - "stream": kwargs["stream"], - "extra_headers": kwargs["extra_headers"] + "model": kwargs.get("model", ""), + "messages": kwargs.get("messages", []), } + + # Safely add optional parameters only if they exist in kwargs + if "temperature" in kwargs: + ollama_supported_params["temperature"] = kwargs["temperature"] if kwargs["temperature"] is not NOT_GIVEN else None + if "top_p" in kwargs: + ollama_supported_params["top_p"] = kwargs["top_p"] if kwargs["top_p"] is not NOT_GIVEN else None + if "max_tokens" in kwargs: + ollama_supported_params["max_tokens"] = kwargs["max_tokens"] if kwargs["max_tokens"] is not NOT_GIVEN else None + + # Add stream parameter - default to False if not present + ollama_supported_params["stream"] = kwargs.get("stream", False) + + # Add extra headers if available + if "extra_headers" in kwargs: + ollama_supported_params["extra_headers"] = kwargs["extra_headers"] + + # IMPORTANT: For tool calls with Ollama (especially with Qwen), + # we need to pass the tools and tool_choice as part of the request + # This is needed for both streaming and non-streaming modes + if "tools" in kwargs and kwargs.get("tools") and kwargs.get("tools") is not NOT_GIVEN: + ollama_supported_params["tools"] = kwargs.get("tools") + + # Include tool_choice if present and not NOT_GIVEN + if "tool_choice" in kwargs and kwargs.get("tool_choice") is not NOT_GIVEN: + ollama_supported_params["tool_choice"] = kwargs.get("tool_choice") # Modify the messages to remove system message for Ollama - if ollama_supported_params["messages"] and ollama_supported_params["messages"][0].get("role") == "system": + if (ollama_supported_params["messages"] and + len(ollama_supported_params["messages"]) > 0 and + ollama_supported_params["messages"][0].get("role") == "system"): # Extract the system message system_content = ollama_supported_params["messages"][0].get("content", "") # Remove it from the messages ollama_supported_params["messages"] = ollama_supported_params["messages"][1:] # If there are user messages, prepend system to first user - if ollama_supported_params["messages"] and ollama_supported_params["messages"][0].get("role") == "user": + if (ollama_supported_params["messages"] and + len(ollama_supported_params["messages"]) > 0 and + ollama_supported_params["messages"][0].get("role") == "user"): # Prepend the system instruction to the first user message, with a separator user_content = ollama_supported_params["messages"][0].get("content", "") if isinstance(user_content, str): @@ -1254,6 +1370,10 @@ class OpenAIChatCompletionsModel(Model): # Remove None values ollama_kwargs = {k: v for k, v in ollama_supported_params.items() if v is not None} + + # Detect if this is a Qwen model + model_str = str(self.model).lower() + is_qwen = "qwen" in model_str if stream: # For streaming with Ollama, we need to create a Response object first @@ -1269,20 +1389,39 @@ class OpenAIChatCompletionsModel(Model): tools=[], parallel_tool_calls=parallel_tool_calls or False, ) - # Get the streaming object - stream_obj = await litellm.acompletion( - **ollama_kwargs, - api_base=get_ollama_api_base().rstrip('/v1'), - custom_llm_provider="ollama" - ) + + # Special handling for Qwen when streaming to ensure tool calls are properly handled + if is_qwen and os.getenv('CAI_STREAM', 'false').lower() == 'true': + # Make sure we have the right custom provider in streaming mode + stream_obj = await litellm.acompletion( + **ollama_kwargs, + api_base=get_ollama_api_base(), + custom_llm_provider="openai" # Use openai provider for best compatibility with Qwen + ) + else: + # Standard Ollama streaming + stream_obj = await litellm.acompletion( + **ollama_kwargs, + api_base=get_ollama_api_base().rstrip('/v1'), + custom_llm_provider="ollama" + ) return response, stream_obj else: # Non-streaming mode - ret = litellm.completion( - **ollama_kwargs, - api_base=get_ollama_api_base().rstrip('/v1'), - custom_llm_provider="ollama" - ) + # For Qwen models, use the openai provider when tools are present + if is_qwen and "tools" in ollama_kwargs: + ret = litellm.completion( + **ollama_kwargs, + api_base=get_ollama_api_base(), + custom_llm_provider="openai" # Use openai provider for better tool handling + ) + else: + # Standard Ollama completion + ret = litellm.completion( + **ollama_kwargs, + api_base=get_ollama_api_base().rstrip('/v1'), + custom_llm_provider="ollama" + ) return ret def _get_client(self) -> AsyncOpenAI: @@ -1290,6 +1429,59 @@ class OpenAIChatCompletionsModel(Model): self._client = AsyncOpenAI() return self._client + # Helper function to detect and format function calls from various models + def _detect_and_format_function_calls(self, delta): + """ + Helper to detect function calls in different formats and normalize them. + Handles Qwen specifics where function calls may be formatted differently. + + Returns: List of normalized tool calls or None + """ + # Standard OpenAI-style tool_calls format + if hasattr(delta, 'tool_calls') and delta.tool_calls: + return delta.tool_calls + elif isinstance(delta, dict) and 'tool_calls' in delta and delta['tool_calls']: + return delta['tool_calls'] + + # Qwen/Ollama function_call format + if isinstance(delta, dict) and 'function_call' in delta: + function_call = delta['function_call'] + return [{ + 'index': 0, + 'id': f"call_{time.time_ns()}", # Generate a unique ID + 'type': 'function', + 'function': { + 'name': function_call.get('name', ''), + 'arguments': function_call.get('arguments', '') + } + }] + + # Anthropic-style tool_use format + if hasattr(delta, 'tool_use') and delta.tool_use: + tool_use = delta.tool_use + return [{ + 'index': 0, + 'id': tool_use.get('id', f"tool_{time.time_ns()}"), + 'type': 'function', + 'function': { + 'name': tool_use.get('name', ''), + 'arguments': tool_use.get('input', '{}') + } + }] + elif isinstance(delta, dict) and 'tool_use' in delta and delta['tool_use']: + tool_use = delta['tool_use'] + return [{ + 'index': 0, + 'id': tool_use.get('id', f"tool_{time.time_ns()}"), + 'type': 'function', + 'function': { + 'name': tool_use.get('name', ''), + 'arguments': tool_use.get('input', '{}') + } + }] + + return None + class _Converter: @classmethod From 584b0c53cdf7d66785f0909625277dfc25436904 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 5 May 2025 16:22:57 +0200 Subject: [PATCH 11/29] fix /agent and /model --- src/cai/cli.py | 33 ++++++++++++++++++++++++++++++++- src/cai/repl/commands/agent.py | 9 +++++---- 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/src/cai/cli.py b/src/cai/cli.py index bdbb54e7..a6e50086 100644 --- a/src/cai/cli.py +++ b/src/cai/cli.py @@ -175,7 +175,8 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= ACTIVE_TIME = 0 idle_time = 0 console = Console() - + last_model = os.getenv('CAI_MODEL', 'qwen2.5:14b') + last_agent_type = os.getenv('CAI_AGENT_TYPE', 'one_tool_agent') # Initialize command completer and key bindings command_completer = FuzzyCommandCompleter() current_text = [''] @@ -208,6 +209,36 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns= while turn_count < max_turns: try: idle_start_time = time.time() + + # Check if model has changed and update if needed + current_model = os.getenv('CAI_MODEL', 'qwen2.5:14b') + if current_model != last_model and hasattr(agent, 'model'): + # Update the model in the agent + if hasattr(agent.model, 'model'): + agent.model.model = current_model + last_model = current_model + + # Check if agent type has changed and recreate agent if needed + current_agent_type = os.getenv('CAI_AGENT_TYPE', 'one_tool_agent') + if current_agent_type != last_agent_type: + try: + # Import is already at the top level + agent = get_agent_by_name(current_agent_type) + last_agent_type = current_agent_type + + # Configure the new agent's model flags + if hasattr(agent, 'model'): + if hasattr(agent.model, 'disable_rich_streaming'): + agent.model.disable_rich_streaming = True + if hasattr(agent.model, 'suppress_final_output'): + agent.model.suppress_final_output = True + + # Apply current model to the new agent + if hasattr(agent.model, 'model'): + agent.model.model = current_model + except Exception as e: + console.print(f"[red]Error switching agent: {str(e)}[/red]") + # Get user input with command completion and history user_input = get_user_input( command_completer, diff --git a/src/cai/repl/commands/agent.py b/src/cai/repl/commands/agent.py index e01d4fdf..dec5b0ee 100644 --- a/src/cai/repl/commands/agent.py +++ b/src/cai/repl/commands/agent.py @@ -191,9 +191,6 @@ class AgentCommand(Command): console.print(table) return True - console.print(table) - return True - def handle_select(self, args: Optional[List[str]] = None) -> bool: # pylint: disable=too-many-branches,line-too-long # noqa: E501 """Handle /agent select command. @@ -226,16 +223,20 @@ class AgentCommand(Command): return False else: # Treat as agent key + selected_agent_key = None for key, agent_obj in agents_to_display.items(): if key == agent_id: agent = agent_obj + selected_agent_key = key agent_name = getattr(agent_obj, "name", key) break else: console.print(f"[red]Error: Unknown agent key: {agent_id}[/red]") return False - os.environ["CAI_MODEL"] = agent_name + # Set the agent key in environment variable (not the agent name) + os.environ["CAI_AGENT_TYPE"] = selected_agent_key + console.print( f"[green]Switched to agent: {agent_name}[/green]") visualize_agent_graph(agent) From f9998768768d5072795a9734e5adaa3ec345a6c6 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Mon, 5 May 2025 16:32:25 +0200 Subject: [PATCH 12/29] change toolbar refresh --- src/cai/repl/ui/toolbar.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/cai/repl/ui/toolbar.py b/src/cai/repl/ui/toolbar.py index 76e9d4e9..44874ac5 100644 --- a/src/cai/repl/ui/toolbar.py +++ b/src/cai/repl/ui/toolbar.py @@ -18,7 +18,7 @@ toolbar_last_refresh = [datetime.datetime.now()] toolbar_cache = { 'html': "", 'last_update': datetime.datetime.now(), - 'refresh_interval': 60 # Refresh every 60 seconds + 'refresh_interval': 5 # Refresh every 60 seconds } # Cache for system information that rarely changes From b5fd3427b086d0e01ec180667e7022660ad5d3b6 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Tue, 6 May 2025 08:50:12 +0200 Subject: [PATCH 13/29] add /history but not tool calls --- .../agents/models/openai_chatcompletions.py | 79 ++++++++++++++----- 1 file changed, 59 insertions(+), 20 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index f1ee8921..0134fcbf 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -102,6 +102,33 @@ _HEADERS = {"User-Agent": _USER_AGENT} message_history = [] +# Function to add a message to history if it's not a duplicate +def add_to_message_history(msg): + """Add a message to history if it's not a duplicate.""" + if not message_history: + message_history.append(msg) + return + + is_duplicate = False + + if msg.get("role") in ["system", "user"]: + is_duplicate = any( + existing.get("role") == msg.get("role") and + existing.get("content") == msg.get("content") + for existing in message_history + ) + + elif msg.get("role") == "assistant" and msg.get("tool_calls"): + is_duplicate = any( + existing.get("role") == "assistant" and + existing.get("tool_calls") and + existing["tool_calls"][0].get("id") == msg["tool_calls"][0].get("id") + for existing in message_history + ) + + if not is_duplicate: + message_history.append(msg) + @dataclass class _StreamingState: started: bool = False @@ -243,25 +270,29 @@ class OpenAIChatCompletionsModel(Model): # --- Add to message_history: user, system, and assistant tool call messages --- # Add system prompt to message_history if system_instructions: - message_history.append({ + sys_msg = { "role": "system", "content": system_instructions - }) + } + add_to_message_history(sys_msg) + # Add user prompt(s) to message_history if isinstance(input, str): - message_history.append({ + user_msg = { "role": "user", "content": input - }) + } + add_to_message_history(user_msg) elif isinstance(input, list): for item in input: # Try to extract user messages if isinstance(item, dict): if item.get("role") == "user": - message_history.append({ + user_msg = { "role": "user", "content": item.get("content", "") - }) + } + add_to_message_history(user_msg) # Get token count estimate before API call for consistent counting estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages) @@ -381,13 +412,15 @@ class OpenAIChatCompletionsModel(Model): } ] } - message_history.append(tool_call_msg) + + add_to_message_history(tool_call_msg) # If the assistant message is just text, add it as well elif hasattr(assistant_msg, "content") and assistant_msg.content: - message_history.append({ + asst_msg = { "role": "assistant", "content": assistant_msg.content - }) + } + add_to_message_history(asst_msg) usage = ( Usage( @@ -458,25 +491,29 @@ class OpenAIChatCompletionsModel(Model): "role": "system", }, ) - # --- Add to message_history: user, system prompts --- + # --- Add to message_history: user, system prompts --- if system_instructions: - message_history.append({ + sys_msg = { "role": "system", "content": system_instructions - }) + } + add_to_message_history(sys_msg) + if isinstance(input, str): - message_history.append({ + user_msg = { "role": "user", "content": input - }) + } + add_to_message_history(user_msg) elif isinstance(input, list): for item in input: if isinstance(item, dict): if item.get("role") == "user": - message_history.append({ + user_msg = { "role": "user", "content": item.get("content", "") - }) + } + add_to_message_history(user_msg) # Get token count estimate before API call for consistent counting estimated_input_tokens, _ = count_tokens_with_tiktoken(converted_messages) @@ -737,6 +774,7 @@ class OpenAIChatCompletionsModel(Model): # Only add if not already in streamed_tool_calls if tool_call_msg not in streamed_tool_calls: streamed_tool_calls.append(tool_call_msg) + add_to_message_history(tool_call_msg) function_call_starting_index = 0 if state.text_content_index_and_output: @@ -940,13 +978,14 @@ class OpenAIChatCompletionsModel(Model): # --- Add assistant tool call(s) to message_history at the end of streaming --- for tool_call_msg in streamed_tool_calls: - message_history.append(tool_call_msg) - # If there was only text output, add that as an assistant message + add_to_message_history(tool_call_msg) + # If there was only text output, add that as an assistant message if (not streamed_tool_calls) and state.text_content_index_and_output and state.text_content_index_and_output[1].text: - message_history.append({ + asst_msg = { "role": "assistant", "content": state.text_content_index_and_output[1].text - }) + } + add_to_message_history(asst_msg) if tracing.include_data(): span_generation.span_data.output = [final_response.model_dump()] From 3568519538bfedb56f0c8eaec0f43f31faa63b54 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Tue, 6 May 2025 08:58:22 +0200 Subject: [PATCH 14/29] print tool calls in /history --- src/cai/repl/commands/history.py | 81 +++++++++++++++++++++++++++++--- 1 file changed, 74 insertions(+), 7 deletions(-) diff --git a/src/cai/repl/commands/history.py b/src/cai/repl/commands/history.py index 38eaeb96..f0610f3a 100644 --- a/src/cai/repl/commands/history.py +++ b/src/cai/repl/commands/history.py @@ -2,8 +2,12 @@ History command for CAI REPL. This module provides commands for displaying conversation history. """ +import json +from typing import Any, Dict, List, Optional + from rich.console import Console # pylint: disable=import-error from rich.table import Table # pylint: disable=import-error +from rich.text import Text # pylint: disable=import-error from cai.repl.commands.base import Command, register_command @@ -20,6 +24,20 @@ class HistoryCommand(Command): description="Display the conversation history", aliases=["/h"] ) + + def handle(self, args: Optional[List[str]] = None, + messages: Optional[List[Dict]] = None) -> bool: + """Handle the history command. + + Args: + args: Optional list of command arguments + messages: Optional list of conversation messages + + Returns: + True if the command was handled successfully, False otherwise + """ + # Currently, the history command doesn't take any arguments + return self.handle_no_args() def handle_no_args(self) -> bool: """Handle the command when no arguments are provided. @@ -27,7 +45,7 @@ class HistoryCommand(Command): Returns: True if the command was handled successfully, False otherwise """ - # Access messages directly from repl.py's global scope + # Access messages directly from openai_chatcompletions.py try: from cai.sdk.agents.models.openai_chatcompletions import message_history # pylint: disable=import-outside-toplevel # noqa: E501 except ImportError: @@ -54,10 +72,11 @@ class HistoryCommand(Command): try: role = msg.get("role", "unknown") content = msg.get("content", "") + tool_calls = msg.get("tool_calls", None) - # Truncate long content for better display - if len(content) > 100: - content = content[:97] + "..." + # Create formatted content based on message type + formatted_content = self._format_message_content( + content, tool_calls) # Color the role based on type if role == "user": @@ -74,14 +93,62 @@ class HistoryCommand(Command): table.add_row( str(idx), f"[{role_style}]{role}[/{role_style}]", - content + formatted_content ) - except Exception: - # Si hay un error, evita esto (no agregues la fila) + except Exception as e: + # Log error but continue with next message + console.print(f"[red]Error displaying message {idx}: {e}[/red]") continue console.print(table) return True + + def _format_message_content( + self, content: Any, tool_calls: List[Dict[str, Any]] + ) -> str: + """Format message content for display, handling both text and tool calls. + + Args: + content: Text content of the message + tool_calls: List of tool calls if present + + Returns: + Formatted string representation of the message content + """ + if tool_calls: + # Format tool calls into a readable string + result = [] + for tc in tool_calls: + func_details = tc.get("function", {}) + func_name = func_details.get("name", "unknown_function") + + # Format arguments (pretty-print JSON if possible) + args_str = func_details.get("arguments", "{}") + try: + # Parse and re-format JSON for better readability + args_dict = json.loads(args_str) + args_formatted = json.dumps(args_dict, indent=2) + # Limit to first 100 chars for display + if len(args_formatted) > 100: + args_formatted = args_formatted[:97] + "..." + except (json.JSONDecodeError, TypeError): + # If not valid JSON, use as is + args_formatted = args_str + if len(args_formatted) > 100: + args_formatted = args_formatted[:97] + "..." + + result.append(f"Function: [bold blue]{func_name}[/bold blue]") + result.append(f"Args: {args_formatted}") + + return "\n".join(result) + elif content: + # Regular text content (truncate if too long) + if len(content) > 100: + return content[:97] + "..." + return content + else: + # No content or tool calls (empty message) + return "[dim italic]Empty message[/dim italic]" # Register the command From cbf5afaa06b636dac3eb45de3d11ede0543b3ada Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Tue, 6 May 2025 10:57:47 +0200 Subject: [PATCH 15/29] solve /flush command --- src/cai/repl/commands/__init__.py | 3 +- src/cai/repl/commands/flush.py | 120 ++++++++++++++++++++++++++++++ 2 files changed, 122 insertions(+), 1 deletion(-) create mode 100644 src/cai/repl/commands/flush.py diff --git a/src/cai/repl/commands/__init__.py b/src/cai/repl/commands/__init__.py index 95bb4a27..cfb0a13a 100644 --- a/src/cai/repl/commands/__init__.py +++ b/src/cai/repl/commands/__init__.py @@ -37,7 +37,8 @@ from cai.repl.commands import ( # pylint: disable=import-error,unused-import,li turns, agent, history, - config + config, + flush ) # Define helper functions diff --git a/src/cai/repl/commands/flush.py b/src/cai/repl/commands/flush.py new file mode 100644 index 00000000..9cee48df --- /dev/null +++ b/src/cai/repl/commands/flush.py @@ -0,0 +1,120 @@ +""" +Flush command for CAI REPL. +This module provides commands for clear the context. +""" +import os +from typing import ( + Dict, + List, + Optional +) +from rich.console import Console # pylint: disable=import-error +from rich.panel import Panel # pylint: disable=import-error +from cai.util import get_model_input_tokens +from cai.repl.commands.base import Command, register_command +from cai.sdk.agents.models.openai_chatcompletions import message_history + +console = Console() + + +class FlushCommand(Command): + """Command to flush the conversation history.""" + + def __init__(self): + """Initialize the flush command.""" + super().__init__( + name="/flush", + description="Clear the current conversation history.", + aliases=["/clear"] + ) + + def handle_no_args(self, messages: Optional[List[Dict]] = None) -> bool: + """Handle the flush command when no args are provided. + + Args: + messages: The conversation history messages + + Returns: + True if the command was handled successfully + """ + # Use both the local messages parameter and the global message_history + local_messages = messages or [] + global_history_length = len(message_history) + + # Get token usage information before clearing + token_info = "" + context_usage = "" + + # Access client through a function to avoid circular imports + # We can use globals() to get the client at runtime + client = self._get_client() + + if client and hasattr(client, 'interaction_input_tokens') and hasattr( + client, 'total_input_tokens'): + model = os.getenv('CAI_MODEL', "qwen2.5:14b") + input_tokens = client.interaction_input_tokens if hasattr( + client, 'interaction_input_tokens') else 0 + total_tokens = client.total_input_tokens if hasattr( + client, 'total_input_tokens') else 0 + max_tokens = get_model_input_tokens(model) + context_pct = (input_tokens / max_tokens) * \ + 100 if max_tokens > 0 else 0 + + token_info = f"Current tokens: {input_tokens}, Total tokens: {total_tokens}" + context_usage = f"Context usage: {context_pct:.1f}% of {max_tokens} tokens" + + # Clear both the local messages list and the global message_history + if local_messages: + local_messages.clear() + + # Always clear the global message history + message_history.clear() + + # Determine which length to report (use the greater of the two) + initial_length = max(len(local_messages) if messages else 0, global_history_length) + + # Display information about the cleared messages + if initial_length > 0: + content = [ + f"Conversation history cleared. Removed {initial_length} messages." + ] + + if token_info: + content.append(token_info) + if context_usage: + content.append(context_usage) + + console.print(Panel( + "\n".join(content), + title="[bold cyan]Context Flushed[/bold cyan]", + border_style="blue", + padding=(1, 2) + )) + else: + console.print(Panel( + "No conversation history to clear.", + title="[bold cyan]Context Flushed[/bold cyan]", + border_style="blue", + padding=(1, 2) + )) + return True + + def _get_client(self): + """Get the CAI client from the global namespace. + + This function avoids circular imports by accessing the client + at runtime instead of import time. + + Returns: + The global CAI client instance or None if not available + """ + try: + # Import here to avoid circular import + from cai.repl.repl import client as global_client # pylint: disable=import-outside-toplevel # noqa: E501 + return global_client + except (ImportError, AttributeError): + return None + + +# Register the /flush command +register_command(FlushCommand()) \ No newline at end of file From d7b8e8b2c849963d9d462dfc97edcf8d7d318d82 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Tue, 6 May 2025 11:07:39 +0200 Subject: [PATCH 16/29] remove turns --- src/cai/repl/commands/__init__.py | 1 - src/cai/repl/commands/turns.py | 96 ------------------------------- 2 files changed, 97 deletions(-) delete mode 100644 src/cai/repl/commands/turns.py diff --git a/src/cai/repl/commands/__init__.py b/src/cai/repl/commands/__init__.py index cfb0a13a..18769f20 100644 --- a/src/cai/repl/commands/__init__.py +++ b/src/cai/repl/commands/__init__.py @@ -34,7 +34,6 @@ from cai.repl.commands import ( # pylint: disable=import-error,unused-import,li platform, kill, model, - turns, agent, history, config, diff --git a/src/cai/repl/commands/turns.py b/src/cai/repl/commands/turns.py deleted file mode 100644 index 3c97f50f..00000000 --- a/src/cai/repl/commands/turns.py +++ /dev/null @@ -1,96 +0,0 @@ -""" -Turns command for CAI REPL. -This module provides commands for viewing and changing the maximum number -of turns. -""" -import os -from typing import ( - List, - Optional -) -from rich.console import Console # pylint: disable=import-error -from rich.panel import Panel # pylint: disable=import-error - -from cai.repl.commands.base import Command, register_command - -console = Console() - - -class TurnsCommand(Command): - """Command for viewing and changing the maximum number of turns.""" - - def __init__(self): - """Initialize the turns command.""" - super().__init__( - name="/turns", - description="View or change the maximum number of turns", - aliases=["/t"] - ) - - def handle(self, args: Optional[List[str]] = None) -> bool: - """Handle the turns command. - - Args: - args: Optional list of command arguments - - Returns: - True if the command was handled successfully, False otherwise - """ - return self.handle_turns_command(args) - - def handle_turns_command(self, args: List[str]) -> bool: - """Change the maximum number of turns for CAI. - - Args: - args: List containing the number of turns - - Returns: - bool: True if the max turns was changed successfully - """ - if not args: - # Display current max turns - max_turns_info = os.getenv("CAI_MAX_TURNS", "inf") - console.print(Panel( - f"Current maximum turns: [bold green]{ - max_turns_info}[/bold green]", - border_style="green", - title="Max Turns Setting" - )) - - # Usage instructions - console.print( - "\n[cyan]Usage:[/cyan] [bold]/turns [/bold]") - console.print("[cyan]Examples:[/cyan]") - console.print(" [bold]/turns 10[/bold] - Limit to 10 turns") - console.print(" [bold]/turns inf[/bold] - Unlimited turns") - return True - - try: - turns = args[0] - # Check if it's a number or 'inf' - if turns.lower() == 'inf': - turns = 'inf' - else: - turns = int(turns) - - # Set the max turns in environment variable - os.environ["CAI_MAX_TURNS"] = turns - - console.print(Panel( - f"Maximum turns changed to: [bold green]{turns}[/bold green]\n" - "[yellow]Note: This will take effect on the next run[/yellow]", - border_style="green", - title="Max Turns Changed" - )) - return True - except ValueError: - console.print(Panel( - "Error: Max turns must be a number or 'inf'", - border_style="red", - title="Invalid Input" - )) - return False - - -# Register the command -register_command(TurnsCommand()) From 66258b0df3b2ceba0ab7584419cbfe52b27dee07 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Tue, 6 May 2025 11:08:27 +0200 Subject: [PATCH 17/29] remove memory --- src/cai/repl/commands/__init__.py | 1 - src/cai/repl/commands/memory.py | 4 ---- 2 files changed, 5 deletions(-) delete mode 100644 src/cai/repl/commands/memory.py diff --git a/src/cai/repl/commands/__init__.py b/src/cai/repl/commands/__init__.py index 18769f20..fee806da 100644 --- a/src/cai/repl/commands/__init__.py +++ b/src/cai/repl/commands/__init__.py @@ -25,7 +25,6 @@ from cai.repl.commands.base import ( # Import all command modules # These imports will register the commands with the registry from cai.repl.commands import ( # pylint: disable=import-error,unused-import,line-too-long,redefined-builtin # noqa: E501,F401 - memory, help, graph, exit, diff --git a/src/cai/repl/commands/memory.py b/src/cai/repl/commands/memory.py deleted file mode 100644 index 12614d92..00000000 --- a/src/cai/repl/commands/memory.py +++ /dev/null @@ -1,4 +0,0 @@ -""" -Memory command for CAI REPL. -This module provides commands for managing memory collections. -""" From 0cfcf2467dff2a58d6fd6f6e21551006d56c0ba2 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Tue, 6 May 2025 12:17:03 +0200 Subject: [PATCH 18/29] add graph capabilities --- src/cai/repl/commands/graph.py | 147 +++++++++++++++++++++++++++++---- 1 file changed, 132 insertions(+), 15 deletions(-) diff --git a/src/cai/repl/commands/graph.py b/src/cai/repl/commands/graph.py index e3ccbc93..d9a32f92 100644 --- a/src/cai/repl/commands/graph.py +++ b/src/cai/repl/commands/graph.py @@ -1,20 +1,64 @@ """ -Graph command for CAI REPL. +Graph command for CAI cli. + This module provides commands for visualizing the agent interaction graph. +It allows users to display a simple directed graph of the conversation history, +showing the sequence of user and agent interactions, including tool calls. """ -from typing import ( - List, - Optional -) +from typing import List, Optional from rich.console import Console # pylint: disable=import-error +from rich.panel import Panel from cai.repl.commands.base import Command, register_command +import os +import importlib.util + console = Console() +def find_agent_name_by_instructions(target_instructions: str, agents_dir: str) -> Optional[str]: + """ + Search all Python files in the agents directory for an agent whose 'instructions' + attribute matches the given target_instructions (ignoring leading/trailing whitespace). + Returns the agent's 'name' attribute if found, otherwise None. + + Args: + target_instructions (str): The instructions string to match. + agents_dir (str): The directory containing agent files. + + Returns: + Optional[str]: The agent name if found, else None. + """ + for filename in os.listdir(agents_dir): + if not filename.endswith(".py") or filename.startswith("__"): + continue + filepath = os.path.join(agents_dir, filename) + try: + spec = importlib.util.spec_from_file_location("agent_mod", filepath) + agent_mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(agent_mod) + for attr_name in dir(agent_mod): + attr = getattr(agent_mod, attr_name) + if hasattr(attr, "instructions"): + agent_instructions = getattr(attr, "instructions", None) + if agent_instructions and agent_instructions.strip() == target_instructions.strip(): + agent_name = getattr(attr, "name", None) + if agent_name: + return agent_name + except Exception: + continue + return None + + class GraphCommand(Command): - """Command for visualizing the agent interaction graph.""" + """ + Command for visualizing the agent interaction graph. + + This command displays a directed graph of the conversation history, + showing the sequence of user and agent messages, and highlighting + tool calls made by the agent. + """ def __init__(self): """Initialize the graph command.""" @@ -25,31 +69,104 @@ class GraphCommand(Command): ) def handle(self, args: Optional[List[str]] = None) -> bool: - """Handle the graph command. + """ + Handle the /graph command. Args: args: Optional list of command arguments Returns: - True if the command was handled successfully, False otherwise + bool: True if the command was handled successfully, False otherwise. """ return self.handle_graph_show() def handle_graph_show(self) -> bool: """Handle /graph show command""" - from cai.repl.repl import client # pylint: disable=import-error - - # Import here to avoid circular imports - - if not client or not client._graph: # pylint: disable=protected-access + from cai.sdk.agents.models.openai_chatcompletions import message_history + if not message_history: console.print("[yellow]No conversation graph available.[/yellow]") return True try: + agents_dir = os.path.join(os.path.dirname(__file__), "../../agents") + agents_dir = os.path.abspath(agents_dir) + + import networkx as nx + + G = nx.DiGraph() + last_agent_name = None + prev_node_idx = None # Track the last node actually added (not system) + + for idx, msg in enumerate(message_history): + role = msg.get("role", "unknown") + # If the message is from the system, update last_agent_name but do not add a node + if role == "system": + system_msg = msg.get("content", "").strip() + agent_name = find_agent_name_by_instructions(system_msg, agents_dir) + if agent_name: + last_agent_name = agent_name + continue + label = role + extra_info = "" + if role == "assistant": + if last_agent_name: + label = last_agent_name + else: + label = "assistant" + if msg.get("tool_calls"): + tool_call = msg["tool_calls"][0] + if tool_call.get("function"): + func_name = tool_call["function"].get("name", "") + func_args = tool_call["function"].get("arguments", "") + extra_info = f"\n[cyan]Tool:[/cyan] [bold]{func_name}[/bold]\n[cyan]Args:[/cyan] {func_args}" + elif role == "user": + user_content = msg.get("content", "") + if user_content: + extra_info = f"\n{user_content}" + label = role + else: + label = role + G.add_node(idx, role=label, extra_info=extra_info) + if prev_node_idx is not None: + G.add_edge(prev_node_idx, idx) + prev_node_idx = idx + + def ascii_graph(G): + """ + Render the conversation graph as a sequence of panels with arrows. + + Args: + G (networkx.DiGraph): The conversation graph. + + Returns: + List: List of rich Panel objects and arrow strings. + """ + lines = [] + node_list = list(G.nodes(data=True)) + for i, (idx, data) in enumerate(node_list): + role = data.get("role", "unknown") + extra_info = data.get("extra_info", "") + role_fmt = f"[bold][blue]{role[:1].upper()}{role[1:]}[/blue][/bold]" + panel_content = f"{role_fmt}" + if extra_info: + panel_content += f"{extra_info}" + panel = Panel( + panel_content, + expand=False, + border_style="cyan" + ) + lines.append(panel) + if i < len(node_list) - 1: + lines.append("[cyan] │\n │\n ▼[/cyan]") + return lines + console.print("\n[bold]Conversation Graph:[/bold]") console.print("------------------") - console.print( - client._graph.ascii()) # pylint: disable=protected-access + if len(G.nodes) == 0: + console.print("[yellow]No messages to display in graph.[/yellow]") + else: + for item in ascii_graph(G): + console.print(item) console.print() return True except Exception as e: # pylint: disable=broad-except From ed68c58d4c30b8c38dc3e7961c44dd46a80055c1 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Tue, 6 May 2025 15:55:21 +0200 Subject: [PATCH 19/29] add workspaces + some virt in common --- src/cai/repl/commands/__init__.py | 3 +- src/cai/repl/commands/config.py | 5 + src/cai/repl/commands/workspace.py | 687 +++++++++++++++++++++++++++++ src/cai/repl/ui/toolbar.py | 15 + src/cai/tools/common.py | 532 ++++++++++++++++++---- 5 files changed, 1143 insertions(+), 99 deletions(-) create mode 100644 src/cai/repl/commands/workspace.py diff --git a/src/cai/repl/commands/__init__.py b/src/cai/repl/commands/__init__.py index fee806da..33e189cb 100644 --- a/src/cai/repl/commands/__init__.py +++ b/src/cai/repl/commands/__init__.py @@ -36,7 +36,8 @@ from cai.repl.commands import ( # pylint: disable=import-error,unused-import,li agent, history, config, - flush + flush, + workspace, ) # Define helper functions diff --git a/src/cai/repl/commands/config.py b/src/cai/repl/commands/config.py index 3588cd19..06a72ebe 100644 --- a/src/cai/repl/commands/config.py +++ b/src/cai/repl/commands/config.py @@ -128,6 +128,11 @@ ENV_VARS = { "description": "Boolean to enable real-time, chunked responses instead of full messages.", "default": "True" }, + 23: { + "name": "CAI_WORKSPACE", + "description": "Name of the current workspace (affects log file naming)", + "default": None + }, } diff --git a/src/cai/repl/commands/workspace.py b/src/cai/repl/commands/workspace.py new file mode 100644 index 00000000..a00ff36c --- /dev/null +++ b/src/cai/repl/commands/workspace.py @@ -0,0 +1,687 @@ +""" +Virtualization command for CAI REPL. +This module provides commands for setting up and managing Docker virtualization +environments. +""" +# Standard library imports +import os +import json +import subprocess +import datetime +import time +from typing import List, Optional, Dict, Any, Tuple + +# Third-party imports +from rich.console import Console +from rich.table import Table +from rich.panel import Panel +from rich.markdown import Markdown +import rich.box + +# Local imports +from cai.repl.commands.base import Command, register_command + +console = Console() + +class WorkspaceCommand(Command): + """Command for workspace management within Docker containers or locally.""" + + def __init__(self): + """Initialize the workspace command.""" + super().__init__( + name="/workspace", + description=( + "Set or display the current workspace name and manage files." + " Affects log file naming and where files are stored." + ), + aliases=["/ws"] + ) + + # Add subcommands + self.add_subcommand( + "set", + "Set the current workspace name", + self.handle_set + ) + self.add_subcommand( + "get", + "Display the current workspace name", + self.handle_get + ) + self.add_subcommand( + "ls", + "List files in the workspace", + self.handle_ls_subcommand + ) + self.add_subcommand( + "exec", + "Execute a command in the workspace", + self.handle_exec_subcommand + ) + self.add_subcommand( + "copy", + "Copy files between host and container", + self.handle_copy_subcommand + ) + + def handle(self, args: Optional[List[str]] = None) -> bool: + """Handle the workspace command. + + Args: + args: Optional list of command arguments + + Returns: + True if the command was handled successfully, False otherwise + """ + # If there are subcommands, process them + if args and args[0] in self.subcommands: + return super().handle(args) + + # No arguments means show workspace info (same as get) + return self.handle_get() + + def handle_no_args(self) -> bool: + """Handle the command when no arguments are provided.""" + return self.handle_get() + + def handle_get(self, _: Optional[List[str]] = None) -> bool: + """Display the current workspace name and directory information.""" + # Get workspace info + workspace_name = os.getenv("CAI_WORKSPACE", None) + + # Check if a container is active + active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") + + # Determine environment (container or host) + if active_container: + try: + # Get container details + result = subprocess.run( + ["docker", "inspect", active_container], + capture_output=True, + text=True, + check=False + ) + + if result.returncode == 0: + container_info = json.loads(result.stdout) + if container_info: + image = container_info[0].get("Config", {}).get("Image", "unknown") + env_type = "container" + env_name = f"Container ({image})" + + # For containers, if workspace is set, use container workspace path + # otherwise use root directory + if workspace_name: + # This will create the workspace in the container if it doesn't exist + workspace_dir = f"/workspace/workspaces/{workspace_name}" + # Ensure the directory exists in the container + subprocess.run( + ["docker", "exec", active_container, "mkdir", "-p", workspace_dir], + capture_output=True, + check=False + ) + else: + workspace_dir = "/" + else: + env_type = "host" + env_name = "Host System (container not running)" + # Use common._get_workspace_dir() for consistency + try: + from cai.tools.common import _get_workspace_dir as get_common_workspace_dir + workspace_dir = get_common_workspace_dir() + except ImportError: + workspace_dir = os.getcwd() # Basic fallback + except Exception: + env_type = "host" + env_name = "Host System (error inspecting container)" + # Use common._get_workspace_dir() for consistency + try: + from cai.tools.common import _get_workspace_dir as get_common_workspace_dir + workspace_dir = get_common_workspace_dir() + except ImportError: + workspace_dir = os.getcwd() # Basic fallback + else: + env_type = "host" + env_name = "Host System" + # Use common._get_workspace_dir() for consistency + try: + from cai.tools.common import _get_workspace_dir as get_common_workspace_dir + workspace_dir = get_common_workspace_dir() + except ImportError: + workspace_dir = os.getcwd() # Basic fallback + + # Show workspace information + console.print( + Panel( + f"Current workspace: [bold green]{workspace_name or 'None'}[/bold green]\n" + f"Working in environment: [bold]{env_name}[/bold]\n" + f"Workspace directory: [bold]{workspace_dir}[/bold]", + title="Workspace Information", + border_style="green" + ) + ) + + # Show available workspace commands + console.print("\n[cyan]Workspace Commands:[/cyan]") + console.print( + " [bold]/workspace set [/bold] - " + "Set the current workspace name") + console.print( + " [bold]/workspace ls[/bold] - " + "List files in the workspace") + console.print( + " [bold]/workspace exec [/bold] - " + "Execute a command in the workspace") + + if active_container: + console.print( + " [bold]/workspace copy [/bold] - " + "Copy files between host and container") + + # List contents of the workspace + self._list_workspace_contents(env_type, workspace_dir) + + return True + + def handle_set(self, args: Optional[List[str]] = None) -> bool: + """Set the current workspace name """ + if not args or len(args) != 1: + console.print( + "[yellow]Usage: /workspace set [/yellow]" + ) + return False + + workspace_name = args[0] + # Allow alphanumeric, underscores, hyphens + if not all(c.isalnum() or c in ['_', '-'] for c in workspace_name): + console.print( + "[red]Invalid workspace name. " + "Use alphanumeric, underscores, or hyphens only.[/red]" + ) + return False + + # Import the necessary modules for setting environment variables + # And for getting workspace dir consistently + try: + from cai.repl.commands.config import set_env_var + from cai.tools.common import _get_workspace_dir as get_common_workspace_dir + from cai.tools.common import _get_container_workspace_path as get_common_container_path + + # Set the environment variable + if not set_env_var("CAI_WORKSPACE", workspace_name): + console.print( + "[red]Failed to set workspace environment variable.[/red]" + ) + return False + except ImportError: + # Fallback if import fails + os.environ["CAI_WORKSPACE"] = workspace_name + # Define basic fallbacks for path functions if import failed + def get_common_workspace_dir(): + base = os.getenv("CAI_WORKSPACE_DIR", ".") # Default to current dir base + name = os.getenv("CAI_WORKSPACE") + if name: + return os.path.abspath(os.path.join(base, name)) + return os.path.abspath(base) # Use base dir if no name + + def get_common_container_path(): + name = os.getenv("CAI_WORKSPACE") + if name: + return f"/workspace/workspaces/{name}" + return "/" # Default container path + + # Get the new workspace directory using the common function + new_workspace_dir = get_common_workspace_dir() + + # Create the directory if it doesn't exist on host + try: # Add try-except for robustness + os.makedirs(new_workspace_dir, exist_ok=True) + except OSError as e: + console.print(f"[red]Error creating host directory {new_workspace_dir}: {e}[/red]") + # Decide if this is fatal or just a warning + + # If container is active, also create the directory in the container + active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") + if active_container: + # Check if container is running + check_process = subprocess.run( + ["docker", "inspect", "--format", "{{.State.Running}}", active_container], + capture_output=True, + text=True, + check=False + ) + + if check_process.returncode == 0 and "true" in check_process.stdout.lower(): + # Get container workspace path using the common function + container_workspace_path = get_common_container_path() + try: + mkdir_cmd = ["docker", "exec", active_container, "mkdir", "-p", container_workspace_path] + mkdir_result = subprocess.run( + mkdir_cmd, + capture_output=True, + text=True, + check=False + ) + + if mkdir_result.returncode == 0: + console.print( + f"[dim]Created workspace directory in container: {container_workspace_path}[/dim]" + ) + else: + console.print( + f"[yellow]Warning: Could not create workspace directory in container: {mkdir_result.stderr}[/yellow]" + ) + except Exception as e: + console.print( + f"[yellow]Warning: Failed to setup workspace in container: {str(e)}[/yellow]" + ) + + # Use a different panel style to indicate success + console.print( + Panel( + f"Workspace changed to: [bold green]{workspace_name}[/bold green]\n" + f"New workspace directory: [bold]{new_workspace_dir}[/bold]", + title="Workspace Updated", + border_style="green" + ) + ) + + return True + + def _get_workspace_dir(self) -> str: + """Get the host workspace directory using the common utility. + + Returns: + The host workspace directory path. + """ + try: + # Use the centralized function from common.py + from cai.tools.common import _get_workspace_dir as get_common_workspace_dir + return get_common_workspace_dir() + except ImportError: + # Provide a basic fallback if import fails, mirroring common.py logic + # without 'cai_default' + base_dir = os.getenv("CAI_WORKSPACE_DIR") + workspace_name = os.getenv("CAI_WORKSPACE") + + if base_dir and workspace_name: + # Basic validation + if not all(c.isalnum() or c in ['_', '-'] for c in workspace_name): + print(f"[yellow]Warning: Invalid CAI_WORKSPACE name '{workspace_name}' in fallback.[/yellow]") + # Fallback to base directory if name is invalid + return os.path.abspath(base_dir) + target_dir = os.path.join(base_dir, workspace_name) + return os.path.abspath(target_dir) + elif base_dir: + # If only base dir is set, use that + return os.path.abspath(base_dir) + else: + # Default to current working directory if nothing else is set + return os.getcwd() + + def _list_workspace_contents(self, env_type: str, workspace_dir: str) -> None: + """List the contents of the workspace. + + Args: + env_type: The environment type (container or host) + workspace_dir: The workspace directory + """ + console.print("\n[bold]Workspace Contents:[/bold]") + + if env_type == "container": + active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") + + # For containers, use the workspace path provided + # This should already be the correct path from handle_get + + # First ensure the workspace directory exists in the container + try: + mkdir_cmd = ["docker", "exec", active_container, "mkdir", "-p", workspace_dir] + subprocess.run( + mkdir_cmd, + capture_output=True, + text=True, + check=False + ) + + # Now list the contents + result = subprocess.run( + ["docker", "exec", active_container, "ls", "-la", workspace_dir], + capture_output=True, + text=True, + check=False + ) + + if result.returncode == 0: + console.print(result.stdout) + else: + console.print(f"[yellow]Error listing container files: {result.stderr}[/yellow]") + # Fallback to host + self._list_host_files(workspace_dir) + except Exception as e: + console.print(f"[yellow]Error accessing container: {str(e)}[/yellow]") + # Fallback to host + self._list_host_files(workspace_dir) + else: + # List files in host + self._list_host_files(workspace_dir) + + def _list_host_files(self, workspace_dir: str) -> None: + """List files in the host workspace. + + Args: + workspace_dir: The workspace directory + """ + # Ensure the directory exists + os.makedirs(workspace_dir, exist_ok=True) + + try: + result = subprocess.run( + ["ls", "-la", workspace_dir], + capture_output=True, + text=True, + check=False + ) + + if result.returncode == 0: + console.print(result.stdout) + else: + console.print(f"[yellow]Error listing files: {result.stderr}[/yellow]") + except Exception as e: + console.print(f"[yellow]Error: {str(e)}[/yellow]") + + def handle_ls_subcommand(self, args: Optional[List[str]] = None) -> bool: + """Handle the ls subcommand. + + Args: + args: Optional list of subcommand arguments + + Returns: + True if the subcommand was handled successfully, False otherwise + """ + # Get workspace info using common functions + try: + from cai.tools.common import _get_workspace_dir as get_common_workspace_dir + from cai.tools.common import _get_container_workspace_path as get_common_container_path + except ImportError: + # Define basic fallbacks if import fails + def get_common_workspace_dir(): + base = os.getenv("CAI_WORKSPACE_DIR", ".") + name = os.getenv("CAI_WORKSPACE") + if name: return os.path.abspath(os.path.join(base, name)) + return os.path.abspath(base) + def get_common_container_path(): + name = os.getenv("CAI_WORKSPACE") + if name: return f"/workspace/workspaces/{name}" + return "/" + + host_workspace_dir = get_common_workspace_dir() + active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") + + # Execute command in the appropriate environment + if active_container: + # Use the container workspace path from common function + container_workspace_path = get_common_container_path() + + # Determine the target path within the container + target_path_in_container = container_workspace_path + if args: + # Ensure args[0] is treated as relative to the workspace + target_path_in_container = os.path.join(container_workspace_path, args[0]) + + # Ensure the base workspace directory exists in the container + mkdir_cmd = ["docker", "exec", active_container, "mkdir", "-p", container_workspace_path] + subprocess.run( + mkdir_cmd, + capture_output=True, + text=True, + check=False + ) + + # Try in container + result = subprocess.run( + ["docker", "exec", active_container, "ls", "-la", target_path_in_container], # Use target path + capture_output=True, + text=True, + check=False + ) + + if result.returncode == 0: + console.print(result.stdout) + return True + + # If failed, try on host + console.print(f"[yellow]Failed to list files in container: {result.stderr}[/yellow]") + console.print("[yellow]Falling back to host system...[/yellow]") + + # List on host + # Determine target path on host relative to host workspace dir + target_path_on_host = host_workspace_dir + if args: + # Ensure args[0] is treated as relative to the workspace + target_path_on_host = os.path.join(host_workspace_dir, args[0]) + + # Ensure the target directory exists on host before listing + # Use os.path.dirname if target is potentially a file path + dir_to_ensure = os.path.dirname(target_path_on_host) if '.' in os.path.basename(target_path_on_host) else target_path_on_host + try: + os.makedirs(dir_to_ensure, exist_ok=True) + except OSError as e: + console.print(f"[red]Error creating directory {dir_to_ensure} on host: {e}[/red]") + # Potentially return False or handle error appropriately + + try: + result = subprocess.run( + ["ls", "-la", target_path_on_host], # Use target path + capture_output=True, + text=True, + check=False + ) + + if result.returncode == 0: + console.print(result.stdout) + return True + else: + console.print(f"[red]Error listing files: {result.stderr}[/red]") + return False + except Exception as e: + console.print(f"[red]Error: {str(e)}[/red]") + return False + + return True + + def handle_exec_subcommand(self, args: Optional[List[str]] = None) -> bool: + """Handle the exec subcommand. + + Args: + args: Optional list of subcommand arguments + + Returns: + True if the subcommand was handled successfully, False otherwise + """ + if not args: + console.print("[yellow]Please specify a command to execute.[/yellow]") + return False + + command = " ".join(args) + # Get workspace info using common functions + try: + from cai.tools.common import _get_workspace_dir as get_common_workspace_dir + from cai.tools.common import _get_container_workspace_path as get_common_container_path + except ImportError: + # Define basic fallbacks if import fails + def get_common_workspace_dir(): + base = os.getenv("CAI_WORKSPACE_DIR", ".") + name = os.getenv("CAI_WORKSPACE") + if name: return os.path.abspath(os.path.join(base, name)) + return os.path.abspath(base) + def get_common_container_path(): + name = os.getenv("CAI_WORKSPACE") + if name: return f"/workspace/workspaces/{name}" + return "/" + + host_workspace_dir = get_common_workspace_dir() + active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") + + # Execute in container if active + if active_container: + try: + # Use the container workspace path from common function + container_workspace_path = get_common_container_path() + + # First ensure the workspace directory exists in the container + mkdir_cmd = ["docker", "exec", active_container, "mkdir", "-p", container_workspace_path] + subprocess.run( + mkdir_cmd, + capture_output=True, + text=True, + check=False + ) + + # Execute the command in the container's workspace directory + result = subprocess.run( + ["docker", "exec", "-w", container_workspace_path, active_container, "sh", "-c", command], + capture_output=True, + text=True, + check=False + ) + + console.print(f"[dim]$ {command}[/dim]") + if result.stdout: + console.print(result.stdout) + + if result.stderr: + console.print(f"[yellow]{result.stderr}[/yellow]") + + if result.returncode != 0: + console.print("[yellow]Command failed in container. Trying on host...[/yellow]") + return self._exec_on_host(command, host_workspace_dir) # Pass host_workspace_dir + + return True + except Exception as e: + console.print(f"[yellow]Error executing in container: {str(e)}[/yellow]") + console.print("[yellow]Falling back to host execution...[/yellow]") + + # Execute on host + return self._exec_on_host(command, host_workspace_dir) # Pass host_workspace_dir + + def _exec_on_host(self, command: str, workspace_dir: str) -> bool: + """Execute a command on the host. + + Args: + command: The command to execute + workspace_dir: The workspace directory + + Returns: + True if the command was executed successfully, False otherwise + """ + # Ensure the directory exists + os.makedirs(workspace_dir, exist_ok=True) + + try: + result = subprocess.run( + command, + shell=True, # nosec B602 + capture_output=True, + text=True, + check=False, + cwd=workspace_dir + ) + + console.print(f"[dim]$ {command}[/dim]") + if result.stdout: + console.print(result.stdout) + + if result.stderr: + console.print(f"[yellow]{result.stderr}[/yellow]") + + return result.returncode == 0 + except Exception as e: + console.print(f"[red]Error executing command: {str(e)}[/red]") + return False + + def handle_copy_subcommand(self, args: Optional[List[str]] = None) -> bool: + """Handle the copy subcommand. + + Args: + args: Optional list of subcommand arguments + + Returns: + True if the subcommand was handled successfully, False otherwise + """ + if not args or len(args) < 2: + console.print("[yellow]Please specify source and destination for copy.[/yellow]") + console.print("Usage: /workspace copy ") + return False + + active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") + if not active_container: + console.print("[yellow]No active container. Copy only works with containers.[/yellow]") + return False + + source = args[0] + destination = args[1] + + # Check if copying from container to host or vice versa + if source.startswith("container:"): + # Copy from container to host + container_path = source[10:] # Remove "container:" prefix + host_path = destination + + if not container_path.startswith("/"): + container_path = f"/workspace/{container_path}" + + try: + result = subprocess.run( + ["docker", "cp", f"{active_container}:{container_path}", host_path], + capture_output=True, + text=True, + check=False + ) + + if result.returncode == 0: + console.print(f"[green]Copied from container:{container_path} to {host_path}[/green]") + return True + else: + console.print(f"[red]Error copying from container: {result.stderr}[/red]") + return False + except Exception as e: + console.print(f"[red]Error: {str(e)}[/red]") + return False + elif destination.startswith("container:"): + # Copy from host to container + host_path = source + container_path = destination[10:] # Remove "container:" prefix + + if not container_path.startswith("/"): + container_path = f"/workspace/{container_path}" + + try: + result = subprocess.run( + ["docker", "cp", host_path, f"{active_container}:{container_path}"], + capture_output=True, + text=True, + check=False + ) + + if result.returncode == 0: + console.print(f"[green]Copied from {host_path} to container:{container_path}[/green]") + return True + else: + console.print(f"[red]Error copying to container: {result.stderr}[/red]") + return False + except Exception as e: + console.print(f"[red]Error: {str(e)}[/red]") + return False + else: + # Ambiguous copy - show help + console.print("[yellow]Ambiguous copy direction. Please specify container: prefix.[/yellow]") + console.print("Examples:") + console.print(" /workspace copy file.txt container:file.txt # Host to container") + console.print(" /workspace copy container:file.txt file.txt # Container to host") + return False + + +# Register the commands +register_command(WorkspaceCommand()) \ No newline at end of file diff --git a/src/cai/repl/ui/toolbar.py b/src/cai/repl/ui/toolbar.py index 44874ac5..95b84728 100644 --- a/src/cai/repl/ui/toolbar.py +++ b/src/cai/repl/ui/toolbar.py @@ -57,7 +57,22 @@ def update_toolbar_in_background(): ip_address = sys_info['ip_address'] os_name = sys_info['os_name'] os_version = sys_info['os_version'] + + # Get the current workspace and base directory + workspace_name = os.getenv("CAI_WORKSPACE") + base_dir = os.getenv("CAI_WORKSPACE_DIR", "workspaces") + # Construct the workspace path + standard_path = os.path.join(base_dir, workspace_name) if workspace_name else "" + workspace_path = "" + if workspace_name: + if os.path.isdir(standard_path): + workspace_path = standard_path + elif os.path.isdir(workspace_name): + workspace_path = os.path.abspath(workspace_name) + else: + workspace_path = standard_path + # Get Ollama information ollama_status = "unavailable" try: diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index e1e58778..e475d5f4 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -24,14 +24,70 @@ except ImportError: # Global dictionary to store active sessions ACTIVE_SESSIONS = {} +def _get_workspace_dir() -> str: + """Determines the target workspace directory based on env vars for host.""" + base_dir_env = os.getenv("CAI_WORKSPACE_DIR") + workspace_name = os.getenv("CAI_WORKSPACE") + + # Determine the base directory + if base_dir_env: + base_dir = os.path.abspath(base_dir_env) + else: # Default base directory is 'workspaces' + if workspace_name: + base_dir = os.path.join(os.getcwd(), "workspaces") + else: # If no workspace name is set, the workspace IS the CWD. + return os.getcwd() + + # If a workspace name is provided, append it to the base directory + if workspace_name: + if not all(c.isalnum() or c in ['_', '-'] for c in workspace_name): + print(color(f"Invalid CAI_WORKSPACE name '{workspace_name}'. " + f"Using directory '{base_dir}' instead.", fg="yellow")) + target_dir = base_dir + else: + target_dir = os.path.join(base_dir, workspace_name) + else: + target_dir = base_dir + + # Ensure the final target directory exists on the host + try: + abs_target_dir = os.path.abspath(target_dir) + os.makedirs(abs_target_dir, exist_ok=True) + return abs_target_dir + except OSError as e: + print(color(f"Error creating/accessing host workspace directory '{abs_target_dir}': {e}", + fg="red")) + print(color(f"Falling back to current directory: {os.getcwd()}", fg="yellow")) + return os.getcwd() + +def _get_container_workspace_path() -> str: + """Determines the target workspace path inside the container.""" + workspace_name = os.getenv("CAI_WORKSPACE") + if workspace_name: + if not all(c.isalnum() or c in ['_', '-'] for c in workspace_name): + print(color(f"Invalid CAI_WORKSPACE name '{workspace_name}' for container. " + f"Using '/workspace'.", fg="yellow")) + return "/" + # Standard path inside CAI containers + return f"/workspace/workspaces/{workspace_name}" + else: + return "/" class ShellSession: # pylint: disable=too-many-instance-attributes """Class to manage interactive shell sessions""" - def __init__(self, command, session_id=None, ctf=None): + def __init__(self, command, session_id=None, ctf=None, workspace_dir=None, container_id=None): # noqa E501 self.session_id = session_id or str(uuid.uuid4())[:8] - self.command = command + self.command = command self.ctf = ctf + self.container_id = container_id + # Determine workspace based on context (container, ctf or local host) + if self.container_id: + self.workspace_dir = _get_container_workspace_path() + elif self.ctf: + self.workspace_dir = workspace_dir or _get_workspace_dir() + else: + self.workspace_dir = _get_workspace_dir() self.process = None self.master = None self.slave = None @@ -40,9 +96,40 @@ class ShellSession: # pylint: disable=too-many-instance-attributes self.last_activity = time.time() def start(self): - """Start the shell session""" + """Start the shell session in the appropriate environment.""" + start_message_cmd = self.command + + # --- Start in Container --- + if self.container_id: + try: + self.master, self.slave = pty.openpty() + docker_cmd_list = [ + "docker", "exec", "-i", + "-w", self.workspace_dir, + self.container_id, + "sh", "-c", # Use shell to handle complex commands if needed + self.command # The actual command to run + ] + self.process = subprocess.Popen( + docker_cmd_list, + stdin=self.slave, + stdout=self.slave, + stderr=self.slave, + preexec_fn=os.setsid, + universal_newlines=True + ) + self.is_running = True + self.output_buffer.append( + f"[Session {self.session_id}] Started in container {self.container_id[:12]}: " + f"{start_message_cmd} in {self.workspace_dir}") + threading.Thread(target=self._read_output, daemon=True).start() + except Exception as e: + self.output_buffer.append(f"Error starting container session: {str(e)}") + self.is_running = False + return + + # --- Start in CTF --- if self.ctf: - # For CTF environments self.is_running = True self.output_buffer.append( f"[Session { @@ -52,81 +139,100 @@ class ShellSession: # pylint: disable=too-many-instance-attributes output = self.ctf.get_shell(self.command) self.output_buffer.append(output) except Exception as e: # pylint: disable=broad-except - self.output_buffer.append(f"Error: {str(e)}") - self.is_running = False + self.output_buffer.append(f"Error executing CTF command: {str(e)}") + self.is_running = False return - # For local environment + # --- Start Locally (Host) --- try: - # Create a pseudo-terminal self.master, self.slave = pty.openpty() - - # Start the process self.process = subprocess.Popen( # pylint: disable=subprocess-popen-preexec-fn, consider-using-with # noqa: E501 - self.command, - shell=True, # nosec B602 + self.command, + shell=True, # nosec B602 stdin=self.slave, stdout=self.slave, stderr=self.slave, - preexec_fn=os.setsid, # Create a new process group + cwd=self.workspace_dir, + preexec_fn=os.setsid, universal_newlines=True ) - self.is_running = True self.output_buffer.append( f"[Session { self.session_id}] Started: { self.command}") - # Start a thread to read output threading.Thread(target=self._read_output, daemon=True).start() except Exception as e: # pylint: disable=broad-except - self.output_buffer.append(f"Error starting session: {str(e)}") + self.output_buffer.append(f"Error starting local session: {str(e)}") self.is_running = False def _read_output(self): """Read output from the process""" try: - while self.is_running: + while self.is_running and self.master is not None: try: + # Check if process has exited before reading + if self.process and self.process.poll() is not None: + self.is_running = False + break + # Read the output output = os.read(self.master, 1024).decode() if output: self.output_buffer.append(output) self.last_activity = time.time() - except OSError: - # No data available or terminal closed - time.sleep(0.1) - if not self.is_process_running(): + else: self.is_running = False break - except Exception as e: # pylint: disable=broad-except - self.output_buffer.append(f"Error reading output: {str(e)}") + except Exception as e: + self.output_buffer.append(f"Error reading output buffer: {str(read_err)}") + self.is_running = False + break + # Add a small sleep to prevent busy-waiting if no output + if is_process_running(self): + time.sleep(0.05) + except Exception as e: + self.output_buffer.append(f"Error in read_output loop: {str(e)}") self.is_running = False + def is_process_running(self): """Check if the process is still running""" + # For CTF or container + if self.container_id or self.ctf: + return self.is_running + # For local host if not self.process: return False return self.process.poll() is None def send_input(self, input_data): - """Send input to the process""" - if not self.is_running: - return "Session is not running" + """Send input to the process (local or container)""" + if not self.is_running: # For CTF or container + if self.process and self.process.poll() is None: + self.is_running = True + else: # For local host + return "Session is not running" try: + # --- Send to CTF --- if self.ctf: - # For CTF environments output = self.ctf.get_shell(input_data) self.output_buffer.append(output) return "Input sent to CTF session" - # For local environment - input_data = input_data.rstrip() + "\n" - os.write(self.master, input_data.encode()) - self.last_activity = time.time() - return "Input sent to session" + # --- Send to Local or Container PTY --- + if self.master is not None: + input_data_bytes = (input_data.rstrip() + "\n").encode() + bytes_written = os.write(self.master, input_data_bytes) + if bytes_written != len(input_data_bytes): + self.output_buffer.append(f"[Session {self.session_id}] Warning: Partial input write.") + self.last_activity = time.time() + return "Input sent to session" + else: + return "Session PTY not available for input" except Exception as e: # pylint: disable=broad-except + self.output_buffer.append(f"Error sending input: {str(e)}") return f"Error sending input: {str(e)}" def get_output(self, clear=True): @@ -138,8 +244,12 @@ class ShellSession: # pylint: disable=too-many-instance-attributes def terminate(self): """Terminate the session""" + session_id_short = self.session_id[:8] if not self.is_running: - return "Session already terminated" + if self.process and self.process.poll() is None: + pass # Process is running, proceed with termination + else: + return f"Session {session_id_short} already terminated or finished." try: self.is_running = False @@ -147,28 +257,64 @@ class ShellSession: # pylint: disable=too-many-instance-attributes if self.process: # Try to terminate the process group try: - os.killpg(os.getpgid(self.process.pid), signal.SIGTERM) - except BaseException: # pylint: disable=bare-except,broad-except # noqa: E501 - # If that fails, try to terminate just the process - self.process.terminate() + os.killpg(os.getpgid(self.process.pid), signal.SIGTERM) + except ProcessLookupError: + pass # Process already gone + except subprocess.TimeoutExpired: + print(color(f"Session {session_id_short} did not terminate gracefully, sending SIGKILL...", fg="yellow")) # noqa E501 + try: + if pgid: + os.killpg(pgid, signal.SIGKILL) # Force kill + else: + self.process.kill() + except ProcessLookupError: + pass # Already gone + except Exception as kill_err: + termination_message = f" (Error during SIGKILL: {kill_err})" + except Exception as term_err: # Catch other errors during SIGTERM + termination_message = f" (Error during SIGTERM: {term_err})" + try: + self.process.kill() + except Exception: pass # Ignore nested errors - # Clean up resources - if self.master: - os.close(self.master) - if self.slave: - os.close(self.slave) - return f"Session {self.session_id} terminated" + # Final check + if self.process.poll() is None: + print(color(f"Session {session_id_short} process {self.process.pid} may still be running after termination attempts.", fg="red")) # noqa E501 + termination_message += " (Warning: Process may still be running)" + + + # Clean up PTY resources if they exist + if self.master: + try: os.close(self.master) + except OSError: pass + self.master = None + if self.slave: + try: os.close(self.slave) + except OSError: pass + self.slave = None + + return termination_message or f"Session {self.session_id} terminated" except Exception as e: # pylint: disable=broad-except - return f"Error terminating session: {str(e)}" + return f"Error terminating session {session_id_short}: {str(e)}" -def create_shell_session(command, ctf=None): - """Create a new shell session""" - session = ShellSession(command, ctf=ctf) +def create_shell_session(command, ctf=None, container_id=None, **kwargs): + """Create a new shell session in the correct workspace/environment.""" + if container_id: + session = ShellSession(command, ctf=ctf, container_id=container_id) + else: + workspace_dir = _get_workspace_dir() + session = ShellSession(command, ctf=ctf, workspace_dir=workspace_dir) + session.start() - ACTIVE_SESSIONS[session.session_id] = session - return session.session_id + if session.is_running or (ctf and not session.is_running): + ACTIVE_SESSIONS[session.session_id] = session + return session.session_id + else: + error_msg = session.get_output(clear=True) + print(color(f"Failed to start session: {error_msg}", fg="red")) + return f"Failed to start session: {error_msg}" def list_shell_sessions(): @@ -212,47 +358,100 @@ def get_session_output(session_id, clear=True): def terminate_session(session_id): """Terminate a specific session""" if session_id not in ACTIVE_SESSIONS: - return f"Session {session_id} not found" + return f"Session {session_id} not found or already terminated." session = ACTIVE_SESSIONS[session_id] result = session.terminate() - del ACTIVE_SESSIONS[session_id] + if session_id in ACTIVE_SESSIONS: + del ACTIVE_SESSIONS[session_id] return result -def _run_ctf(ctf, command, stdout=False, timeout=100, stream=False, call_id=None): +def _run_ctf(ctf, command, stdout=False, timeout=100, workspace_dir=None): + """Runs command in CTF env, changing to workspace_dir first.""" + target_dir = workspace_dir or _get_workspace_dir() + full_command = f"cd '{target_dir}' && {command}" + original_cmd_for_msg = command # For logging + context_msg = f"(ctf:{target_dir})" try: - # Ensure the command is executed in a shell that supports command - # chaining - output = ctf.get_shell(command, timeout=timeout) - # exploit_logger.log_ok() - + output = ctf.get_shell(full_command, timeout=timeout) if stdout: - print("\033[32m" + output + "\033[0m") - return output # output if output else result.stder + print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501 + return output except Exception as e: # pylint: disable=broad-except - print(color(f"Error executing CTF command: {e}", fg="red")) - # exploit_logger.log_error(str(e)) - return f"Error executing CTF command: {str(e)}" + error_msg = f"Error executing CTF command '{original_cmd_for_msg}' in '{target_dir}': {e}" # noqa E501 + print(color(error_msg, fg="red")) + return error_msg + +def _run_ssh(command, stdout=False, timeout=100, workspace_dir=None): + """Runs command via SSH. Assumes SSH agent or passwordless setup unless sshpass is used externally.""" # noqa E501 + ssh_user = os.environ.get('SSH_USER') + ssh_host = os.environ.get('SSH_HOST') + ssh_pass = os.environ.get('SSH_PASS') + remote_command = command + original_cmd_for_msg = command + context_msg = f"({ssh_user}@{ssh_host})" + + # Construct base SSH command list + if ssh_pass: + ssh_cmd_list = ["sshpass", "-p", ssh_pass, "ssh", f"{ssh_user}@{ssh_host}"] # noqa E501 + else: + ssh_cmd_list = ["ssh", f"{ssh_user}@{ssh_host}"] + ssh_cmd_list.append(remote_command) + + try: + # Use subprocess.run with list of args for better security than shell=True + result = subprocess.run( + ssh_cmd_list, + capture_output=True, + text=True, + check=False, # Don't raise exception on non-zero exit code + timeout=timeout + ) + output = result.stdout if result.stdout else result.stderr + if stdout: + print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501 + # Return combined output, potentially including errors + return output.strip() + except subprocess.TimeoutExpired as e: + error_output = e.stdout if e.stdout else str(e) + timeout_msg = f"Timeout executing SSH command: {error_output}" + if stdout: + print(f"\033[33m{context_msg} $ {original_cmd_for_msg}\nTIMEOUT\n{error_output}\033[0m") # noqa E501 + return timeout_msg + except FileNotFoundError: + # Handle case where ssh or sshpass isn't installed + error_msg = f"'sshpass' or 'ssh' command not found. Ensure they are installed and in PATH." # noqa E501 + print(color(error_msg, fg="red")) + return error_msg + except Exception as e: # pylint: disable=broad-except + error_msg = f"Error executing SSH command '{original_cmd_for_msg}' on {ssh_host}: {e}" # noqa E501 + print(color(error_msg, fg="red")) + return error_msg -def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None): +def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None, workspace_dir=None): + """Runs command locally in the specified workspace_dir.""" # If streaming is enabled and we have a call_id if stream and call_id: - return _run_local_streamed(command, call_id, timeout, tool_name) + return _run_local_streamed(command, call_id, timeout, tool_name, workspace_dir) + target_dir = workspace_dir or _get_workspace_dir() + original_cmd_for_msg = command # For logging + context_msg = f"(local:{target_dir})" try: - # nosec B602 - shell=True is required for command chaining result = subprocess.run( command, shell=True, # nosec B602 capture_output=True, text=True, - check=False, - timeout=timeout) + check=False, + timeout=timeout, + cwd=target_dir + ) output = result.stdout if result.stdout else result.stderr if stdout: - print("\033[32m" + output + "\033[0m") + print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501 # Skip passing output to cli_print_tool_output when CAI_STREAM=true # This prevents duplicate output in streaming mode @@ -261,20 +460,21 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t # Optional: Add cli_print_tool_output call here if needed for non-streaming pass - return output + return output.strip() except subprocess.TimeoutExpired as e: - error_output = e.stdout.decode() if e.stdout else str(e) + error_output = e.stdout if e.stdout else str(e) if stdout: print("\033[32m" + error_output + "\033[0m") - return error_output + return error_output except Exception as e: # pylint: disable=broad-except - error_msg = f"Error executing local command: {e}" - print(color(error_msg, fg="red")) - return error_msg + error_msg = f"Error executing local command: {e}" + print(color(error_msg, fg="red")) + return error_msg -def _run_local_streamed(command, call_id, timeout=100, tool_name=None): - """Run a local command with streaming output to the Tool output panel""" +def _run_local_streamed(command, call_id, timeout=100, tool_name=None, workspace_dir=None): + """Run a local command with streaming output to the Tool output panel.""" + target_dir = workspace_dir or _get_workspace_dir() try: # Try to import Rich for nice display try: @@ -298,7 +498,8 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, - bufsize=1 + bufsize=1, + cwd=target_dir # Set CWD for local process ) # If tool_name is not provided, derive it from the command @@ -328,7 +529,7 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): header.append(")", style="yellow") tool_time = 0 start_time = time.time() - total_time = time.time() - START_TIME + total_time = time.time() - START_TIME timing_info = [] if total_time: timing_info.append(f"Total: {format_time(total_time)}") @@ -505,7 +706,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg async_mode=False, session_id=None, timeout=100, stream=False, call_id=None, tool_name=None): """ - Run command either in CTF container or on the local attacker machine + Run command in the appropriate environment (Docker, CTF, SSH, Local) + and workspace. Args: command: The command to execute @@ -520,34 +722,168 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg If None, the tool name will be derived from the command. Returns: - str: Command output, status message, or session ID + str: Command output, status message, or session ID. """ # If session_id is provided, send command to that session if session_id: if session_id not in ACTIVE_SESSIONS: return f"Session {session_id} not found" - - result = send_to_session(session_id, command) + session = ACTIVE_SESSIONS[session_id] + result = session.send_input(command) # Send the raw command string if stdout: output = get_session_output(session_id, clear=False) - print("\033[32m" + output + "\033[0m") - return result - - # If async_mode, create a new session - if async_mode: - session_id = create_shell_session(command, ctf) - if stdout: - # Wait a moment for initial output - 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." + env_type = "Local" + if session.container_id: + env_type = f"Container({session.container_id[:12]})" + elif session.ctf: + env_type = "CTF" + print(f"\033[32m(Session {session_id} in {env_type}:{session.workspace_dir}) >> {command}\n{output}\033[0m") # noqa E501 + return result # Return the result of sending input ("Input sent..." or error) # 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 + + # 2. Determine Execution Environment (Container > CTF > SSH > Local) + active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") + is_ssh_env = all(os.getenv(var) for var in ['SSH_USER', 'SSH_HOST']) + + # --- Docker Container Execution --- + if active_container and not ctf and not is_ssh_env: + container_id = active_container + container_workspace = _get_container_workspace_path() + context_msg = f"(docker:{container_id[:12]}:{container_workspace})" + + # Handle Async Session Creation in Container + if async_mode: + # Create a session specifically for the container environment + new_session_id = create_shell_session(command, container_id=container_id) # noqa E501 + if "Failed" in new_session_id: # Check if session creation failed + return new_session_id + if stdout: + # Wait a moment for initial output + time.sleep(0.2) + output = get_session_output(new_session_id, clear=False) + print(f"\033[32m(Started Session {new_session_id} in {context_msg})\n{output}\033[0m") # noqa E501 + return f"Started async session {new_session_id} in container {container_id[:12]}. Use this ID to interact." # noqa E501 + + # Handle Streaming Container Execution - not yet implemented for containers + if stream: + # For now, display that streaming isn't supported for containers + from cai.util import cli_print_tool_output + if call_id and tool_name: + tool_args = {"command": command, "container": container_id[:12]} + cli_print_tool_output( + tool_name, + tool_args, + "Streaming not yet supported for container execution. Running normally...", + call_id=call_id + ) + + # Handle Synchronous Execution in Container + try: + # Ensure container workspace exists (best effort) + # Consider moving this to workspace set/container activation + mkdir_cmd = ["docker", "exec", container_id, "mkdir", "-p", container_workspace] # noqa E501 + subprocess.run(mkdir_cmd, capture_output=True, text=True, check=False, timeout=10) # noqa E501 + + # Construct the docker exec command with workspace context + cmd_list = [ + "docker", "exec", + "-w", container_workspace, # Set working directory + container_id, + "sh", "-c", command # Execute command via shell + ] + result = subprocess.run( + cmd_list, + capture_output=True, + text=True, + check=False, # Don't raise exception on non-zero exit + timeout=timeout + ) + + output = result.stdout if result.stdout else result.stderr + output = output.strip() # Clean trailing newline + + if stdout: + print(f"\033[32m{context_msg} $ {command}\n{output}\033[0m") # noqa E501 + + # Check if command failed specifically because container isn't running + if result.returncode != 0 and "is not running" in result.stderr: + print(color(f"{context_msg} Container is not running. Attempting execution on host instead.", fg="yellow")) # noqa E501 + # Fallback to local execution, preserving workspace context + return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 + + return output # Return combined stdout/stderr + + except subprocess.TimeoutExpired: + timeout_msg = "Timeout executing command in container." + if stdout: + print(f"\033[33m{context_msg} $ {command}\nTIMEOUT\033[0m") # noqa E501 + print(color("Attempting execution on host instead.", fg="yellow")) + # Fallback to local execution on timeout + return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 + except Exception as e: # pylint: disable=broad-except + error_msg = f"Error executing command in container: {str(e)}" + print(color(f"{context_msg} {error_msg}", fg="red")) + print(color("Attempting execution on host instead.", fg="yellow")) + # Fallback to local execution on other errors + return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 + + # --- CTF Execution --- if ctf: - return _run_ctf(ctf, command, stdout, timeout, stream, call_id) - return _run_local(command, stdout, timeout, stream, call_id, tool_name) + # Handling streaming for CTF - not fully implemented yet + if stream: + from cai.util import cli_print_tool_output + if call_id and tool_name: + tool_args = {"command": command, "ctf": True} + cli_print_tool_output( + tool_name, + tool_args, + "Streaming not yet supported for CTF execution. Running normally...", + call_id=call_id + ) + + # _run_ctf handles workspace internally using _get_workspace_dir() default + return _run_ctf(ctf, command, stdout, timeout) # Pass None for workspace_dir + + # --- SSH Execution --- + if is_ssh_env: + # Async for SSH would require session management via SSH client features + if async_mode: + return "Async mode not fully supported for SSH environment via this function yet." + + # Handling streaming for SSH - not fully implemented yet + if stream: + from cai.util import cli_print_tool_output + if call_id and tool_name: + tool_args = {"command": command, "ssh": True} + cli_print_tool_output( + tool_name, + tool_args, + "Streaming not yet supported for SSH execution. Running normally...", + call_id=call_id + ) + + # _run_ssh handles command execution, workspace is relative to remote home + return _run_ssh(command, stdout, timeout) # Workspace dir less relevant here + + # --- Local Execution (Default Fallback) --- + # Let _run_local handle determining the host workspace + # Handle Async Session Creation Locally + if async_mode: + # create_shell_session uses _get_workspace_dir() when container_id is None + new_session_id = create_shell_session(command) + if isinstance(new_session_id, str) and "Failed" in new_session_id: # Check failure + return new_session_id + # Retrieve the actual workspace dir the session is using + session = ACTIVE_SESSIONS.get(new_session_id) + actual_workspace = session.workspace_dir if session else "unknown" + if stdout: + time.sleep(0.2) # Allow session buffer to populate + output = get_session_output(new_session_id, clear=False) + print(f"\033[32m(Started Session {new_session_id} in local:{actual_workspace})\n{output}\033[0m") + return f"Started async session {new_session_id} locally. Use this ID to interact." + + # Handle Synchronous Execution Locally using _run_local default with streaming support + return _run_local(command, stdout, timeout, stream, call_id, tool_name, None) From 78084110188677b0180d7d4ac37d595ea82d3ed6 Mon Sep 17 00:00:00 2001 From: luijait Date: Tue, 6 May 2025 16:21:02 +0200 Subject: [PATCH 20/29] Qwen executes tool calls --- .../agents/models/openai_chatcompletions.py | 485 +++++++++++++++--- src/cai/sdk/agents/run.py | 81 +++ src/cai/tools/common.py | 99 +++- 3 files changed, 576 insertions(+), 89 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 0134fcbf..c5138d14 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -8,6 +8,8 @@ import litellm import tiktoken import inspect import hashlib +import re +import asyncio from collections.abc import AsyncIterator, Iterable from dataclasses import dataclass, field @@ -540,7 +542,27 @@ class OpenAIChatCompletionsModel(Model): streaming_text_buffer = "" # For tool call streaming, accumulate tool_calls to add to message_history at the end streamed_tool_calls = [] + + # Ollama specific: accumulate full content to check for function calls at the end + # Some Ollama models output the function call as JSON in the text content + ollama_full_content = "" + is_ollama = False + + model_str = str(self.model).lower() + is_ollama = self.is_ollama or "ollama" in model_str or ":" in model_str or "qwen" in model_str + # Add a small delay to make sure any previous tool outputs are fully rendered + # This helps prevent overlapping of panels in the terminal + await asyncio.sleep(0.2) + + # Add visual separation before agent output + if streaming_context and should_show_rich_stream: + # If we're using rich context, we'll add separation through that + pass + else: + # Print clear visual separator + print("\n") + async for chunk in stream: if not state.started: state.started = True @@ -593,6 +615,10 @@ class OpenAIChatCompletionsModel(Model): content = delta['content'] if content: + # For Ollama, we need to accumulate the full content to check for function calls + if is_ollama: + ollama_full_content += content + # Add to the streaming text buffer streaming_text_buffer += content @@ -776,6 +802,320 @@ class OpenAIChatCompletionsModel(Model): streamed_tool_calls.append(tool_call_msg) add_to_message_history(tool_call_msg) + # Special handling for Ollama - check if accumulated text contains a valid function call + if is_ollama and ollama_full_content and len(state.function_calls) == 0: + # Look for JSON object that might be a function call + try: + # Try to extract a JSON object from the content + json_start = ollama_full_content.find('{') + json_end = ollama_full_content.rfind('}') + 1 + + logger.debug(f"Ollama content length: {len(ollama_full_content)}, JSON start: {json_start}, JSON end: {json_end}") + + if json_start >= 0 and json_end > json_start: + json_str = ollama_full_content[json_start:json_end] + logger.debug(f"Extracted potential JSON: {json_str[:100]}...") + + # Special check for generic_linux_command format + if "generic_linux_command" in json_str and "arguments" in json_str: + logger.debug("Detected generic_linux_command pattern - using special handling") + + # Try regex pattern matching to handle potentially malformed JSON + cmd_pattern = re.search(r'"command"\s*:\s*"([^"]+)"', json_str) + args_pattern = re.search(r'"args"\s*:\s*"([^"]*)"', json_str) + ctf_pattern = re.search(r'"ctf"\s*:\s*"([^"]*)"', json_str) + + if cmd_pattern: + # Create a tool call specifically for generic_linux_command + command = cmd_pattern.group(1) + args = args_pattern.group(1) if args_pattern else "" + ctf = ctf_pattern.group(1) if ctf_pattern else "" + + tool_call_id = f"call_{hashlib.md5(('generic_linux_command' + str(time.time())).encode()).hexdigest()[:8]}" + + # Create a proper arguments object + command_args = { + "command": command, + "args": args, + "ctf": ctf, + "async_mode": False, + "session_id": "" + } + + arguments_str = json.dumps(command_args) + logger.debug(f"Created generic_linux_command with args: {arguments_str}") + + # Add it to our function_calls state + state.function_calls[0] = ResponseFunctionToolCall( + id=FAKE_RESPONSES_ID, + arguments=arguments_str, + name="generic_linux_command", + type="function_call", + call_id=tool_call_id, + ) + + # Display the tool call in CLI + from cai.util import cli_print_agent_messages + try: + # Create a message-like object to display the function call + tool_msg = type('ToolCallWrapper', (), { + 'content': None, + 'tool_calls': [ + type('ToolCallDetail', (), { + 'function': type('FunctionDetail', (), { + 'name': "generic_linux_command", + 'arguments': arguments_str + }), + 'id': tool_call_id, + 'type': 'function' + }) + ] + }) + + # Print the tool call using the CLI utility + cli_print_agent_messages( + agent_name=getattr(self, 'agent_name', 'Agent'), + message=tool_msg, + counter=getattr(self, 'interaction_counter', 0), + model=str(self.model), + debug=False, + interaction_input_tokens=estimated_input_tokens, + interaction_output_tokens=estimated_output_tokens, + interaction_reasoning_tokens=0, + total_input_tokens=getattr(self, 'total_input_tokens', 0) + estimated_input_tokens, + total_output_tokens=getattr(self, 'total_output_tokens', 0) + estimated_output_tokens, + total_reasoning_tokens=0, + interaction_cost=None, + total_cost=None + ) + except Exception as e: + logger.error(f"Error displaying tool call in CLI: {e}") + + # Add to message history + tool_call_msg = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": tool_call_id, + "type": "function", + "function": { + "name": "generic_linux_command", + "arguments": arguments_str + } + } + ] + } + + streamed_tool_calls.append(tool_call_msg) + add_to_message_history(tool_call_msg) + logger.debug(f"Added generic_linux_command with args: {arguments_str}") + + + # Standard JSON parsing for other function calls + parsed = json.loads(json_str) + + # Check if it looks like a function call + if ('name' in parsed and 'arguments' in parsed and + isinstance(parsed['arguments'], dict)): + + logger.debug(f"Found valid function call in Ollama output: {json_str}") + + # Create a tool call from the JSON object + tool_call_id = f"call_{hashlib.md5((parsed['name'] + str(time.time())).encode()).hexdigest()[:8]}" + + # Add it to our function_calls state + state.function_calls[0] = ResponseFunctionToolCall( + id=FAKE_RESPONSES_ID, + arguments=json.dumps(parsed['arguments']), + name=parsed['name'], + type="function_call", + call_id=tool_call_id, + ) + + # Ensure arguments is a valid JSON string + arguments_str = "" + if isinstance(parsed['arguments'], dict): + arguments_str = json.dumps(parsed['arguments']) + elif isinstance(parsed['arguments'], str): + # If it's already a string, check if it's valid JSON + try: + # Try parsing to validate + json.loads(parsed['arguments']) + arguments_str = parsed['arguments'] + except: + # If not valid JSON, encode it as a JSON string + arguments_str = json.dumps(parsed['arguments']) + else: + # For any other type, convert to string and then JSON + arguments_str = json.dumps(str(parsed['arguments'])) + + logger.debug(f"Final arguments string: {arguments_str}") + + # Update with the properly formatted arguments + state.function_calls[0].arguments = arguments_str + + # Log the tool call for development purposes + logger.debug(f"Adding tool call: {parsed['name']} with arguments: {arguments_str}") + + # Display the tool call in CLI + from cai.util import cli_print_agent_messages + try: + # Create a message-like object to display the function call + tool_msg = type('ToolCallWrapper', (), { + 'content': None, + 'tool_calls': [ + type('ToolCallDetail', (), { + 'function': type('FunctionDetail', (), { + 'name': parsed['name'], + 'arguments': arguments_str + }), + 'id': tool_call_id, + 'type': 'function' + }) + ] + }) + + # Print the tool call using the CLI utility + cli_print_agent_messages( + agent_name=getattr(self, 'agent_name', 'Agent'), + message=tool_msg, + counter=getattr(self, 'interaction_counter', 0), + model=str(self.model), + debug=False, + interaction_input_tokens=estimated_input_tokens, + interaction_output_tokens=estimated_output_tokens, + interaction_reasoning_tokens=0, # Not available for Ollama + total_input_tokens=getattr(self, 'total_input_tokens', 0) + estimated_input_tokens, + total_output_tokens=getattr(self, 'total_output_tokens', 0) + estimated_output_tokens, + total_reasoning_tokens=getattr(self, 'total_reasoning_tokens', 0), + interaction_cost=None, + total_cost=None, + tool_output=None # Will be shown once the tool is executed + ) + except Exception as e: + logger.error(f"Error displaying tool call in CLI: {e}") + + # Add to message history + tool_call_msg = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": tool_call_id, + "type": "function", + "function": { + "name": parsed['name'], + "arguments": arguments_str + } + } + ] + } + + streamed_tool_calls.append(tool_call_msg) + add_to_message_history(tool_call_msg) + + logger.debug(f"Added function call: {parsed['name']} with args: {json.dumps(parsed['arguments'])}") + except Exception as e: + logger.debug(f"Failed to parse potential Ollama function call: {e}") + + # Even if JSON parsing fails, try to extract generic_linux_command with regex + if "generic_linux_command" in ollama_full_content: + logger.debug("JSON parsing failed but detected generic_linux_command - trying regex fallback") + try: + # Use regex to extract command components even from malformed output + cmd_pattern = re.search(r'"command"\s*:\s*"([^"]+)"', ollama_full_content) + args_pattern = re.search(r'"args"\s*:\s*"([^"]*)"', ollama_full_content) + ctf_pattern = re.search(r'"ctf"\s*:\s*"([^"]*)"', ollama_full_content) + + if cmd_pattern: + # Create a tool call specifically for generic_linux_command + command = cmd_pattern.group(1) + args = args_pattern.group(1) if args_pattern else "" + ctf = ctf_pattern.group(1) if ctf_pattern else "" + + tool_call_id = f"call_{hashlib.md5(('generic_linux_command' + str(time.time())).encode()).hexdigest()[:8]}" + + # Create a proper arguments object + command_args = { + "command": command, + "args": args, + "ctf": ctf, + "async_mode": False, + "session_id": "" + } + + arguments_str = json.dumps(command_args) + logger.debug(f"Fallback created generic_linux_command with args: {arguments_str}") + + # Add it to our function_calls state + state.function_calls[0] = ResponseFunctionToolCall( + id=FAKE_RESPONSES_ID, + arguments=arguments_str, + name="generic_linux_command", + type="function_call", + call_id=tool_call_id, + ) + + # Display the tool call in CLI + from cai.util import cli_print_agent_messages + try: + # Create a message-like object for display + tool_msg = type('ToolCallWrapper', (), { + 'content': None, + 'tool_calls': [ + type('ToolCallDetail', (), { + 'function': type('FunctionDetail', (), { + 'name': "generic_linux_command", + 'arguments': arguments_str + }), + 'id': tool_call_id, + 'type': 'function' + }) + ] + }) + + cli_print_agent_messages( + agent_name=getattr(self, 'agent_name', 'Agent'), + message=tool_msg, + counter=getattr(self, 'interaction_counter', 0), + model=str(self.model), + debug=False, + interaction_input_tokens=estimated_input_tokens, + interaction_output_tokens=estimated_output_tokens, + interaction_reasoning_tokens=0, + total_input_tokens=getattr(self, 'total_input_tokens', 0) + estimated_input_tokens, + total_output_tokens=getattr(self, 'total_output_tokens', 0) + estimated_output_tokens, + total_reasoning_tokens=0, + interaction_cost=None, + total_cost=None + ) + except Exception as e: + logger.error(f"Error displaying fallback tool call in CLI: {e}") + + # Add to message history + tool_call_msg = { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": tool_call_id, + "type": "function", + "function": { + "name": "generic_linux_command", + "arguments": arguments_str + } + } + ] + } + + streamed_tool_calls.append(tool_call_msg) + add_to_message_history(tool_call_msg) + + logger.debug(f"Added fallback generic_linux_command") + except Exception as regex_err: + logger.error(f"Regex fallback also failed: {regex_err}") + function_call_starting_index = 0 if state.text_content_index_and_output: function_call_starting_index += 1 @@ -952,13 +1292,26 @@ class OpenAIChatCompletionsModel(Model): direct_stats = final_stats.copy() direct_stats["interaction_cost"] = float(interaction_cost) direct_stats["total_cost"] = float(total_cost) + + # Add a small delay to avoid overlapping outputs + await asyncio.sleep(0.3) + # Use the direct copy with guaranteed float costs finish_agent_streaming(streaming_context, direct_stats) + + # Add visual separation after agent output completes + print("\n") # If we're not using rich streaming and not suppressing output, use old method elif not self.suppress_final_output and final_response.output and any(isinstance(item, ResponseOutputMessage) for item in final_response.output): # Find the assistant message to print for item in final_response.output: if isinstance(item, ResponseOutputMessage) and item.role == 'assistant': + # Add a small delay to avoid overlapping outputs + await asyncio.sleep(0.3) + + # Print clear visual separator before message + print("\n") + cli_print_agent_messages( agent_name=getattr(self, 'agent_name', 'Agent'), message=item, @@ -974,6 +1327,9 @@ class OpenAIChatCompletionsModel(Model): interaction_cost=interaction_cost, total_cost=total_cost, ) + + # Add visual separation after message + print("\n") break # --- Add assistant tool call(s) to message_history at the end of streaming --- @@ -1194,6 +1550,7 @@ class OpenAIChatCompletionsModel(Model): # Use the specialized Qwen approach first return await self._fetch_response_litellm_ollama(kwargs, model_settings, tool_choice, stream, parallel_tool_calls) except Exception as qwen_e: + print(qwen_e) # If that fails, try our direct OpenAI approach qwen_params = kwargs.copy() qwen_params["api_base"] = get_ollama_api_base() @@ -1351,117 +1708,80 @@ class OpenAIChatCompletionsModel(Model): # Standard OpenAI handling for non-streaming ret = litellm.completion(**kwargs) return ret - + async def _fetch_response_litellm_ollama( self, kwargs: dict, model_settings: ModelSettings, tool_choice: ChatCompletionToolChoiceOptionParam | NotGiven, stream: bool, - parallel_tool_calls: bool + parallel_tool_calls: bool, + provider="ollama" ) -> ChatCompletion | tuple[Response, AsyncStream[ChatCompletionChunk]]: + # Extract only supported parameters for Ollama ollama_supported_params = { "model": kwargs.get("model", ""), "messages": kwargs.get("messages", []), + "stream": kwargs.get("stream", False) } - # Safely add optional parameters only if they exist in kwargs - if "temperature" in kwargs: - ollama_supported_params["temperature"] = kwargs["temperature"] if kwargs["temperature"] is not NOT_GIVEN else None - if "top_p" in kwargs: - ollama_supported_params["top_p"] = kwargs["top_p"] if kwargs["top_p"] is not NOT_GIVEN else None - if "max_tokens" in kwargs: - ollama_supported_params["max_tokens"] = kwargs["max_tokens"] if kwargs["max_tokens"] is not NOT_GIVEN else None - - # Add stream parameter - default to False if not present - ollama_supported_params["stream"] = kwargs.get("stream", False) + # Add optional parameters if they exist and are not NOT_GIVEN + for param in ["temperature", "top_p", "max_tokens"]: + if param in kwargs and kwargs[param] is not NOT_GIVEN: + ollama_supported_params[param] = kwargs[param] # Add extra headers if available if "extra_headers" in kwargs: ollama_supported_params["extra_headers"] = kwargs["extra_headers"] - # IMPORTANT: For tool calls with Ollama (especially with Qwen), - # we need to pass the tools and tool_choice as part of the request - # This is needed for both streaming and non-streaming modes + # Add tools and tool_choice for compatibility with Qwen if "tools" in kwargs and kwargs.get("tools") and kwargs.get("tools") is not NOT_GIVEN: ollama_supported_params["tools"] = kwargs.get("tools") - # Include tool_choice if present and not NOT_GIVEN if "tool_choice" in kwargs and kwargs.get("tool_choice") is not NOT_GIVEN: ollama_supported_params["tool_choice"] = kwargs.get("tool_choice") - # Modify the messages to remove system message for Ollama - if (ollama_supported_params["messages"] and - len(ollama_supported_params["messages"]) > 0 and - ollama_supported_params["messages"][0].get("role") == "system"): - # Extract the system message - system_content = ollama_supported_params["messages"][0].get("content", "") - # Remove it from the messages - ollama_supported_params["messages"] = ollama_supported_params["messages"][1:] - # If there are user messages, prepend system to first user - if (ollama_supported_params["messages"] and - len(ollama_supported_params["messages"]) > 0 and - ollama_supported_params["messages"][0].get("role") == "user"): - # Prepend the system instruction to the first user message, with a separator - user_content = ollama_supported_params["messages"][0].get("content", "") - if isinstance(user_content, str): - ollama_supported_params["messages"][0]["content"] = f"System: {system_content}\n\nUser: {user_content}" - # Remove None values ollama_kwargs = {k: v for k, v in ollama_supported_params.items() if v is not None} - # Detect if this is a Qwen model + # Check if this is a Qwen model model_str = str(self.model).lower() is_qwen = "qwen" in model_str - + + api_base = get_ollama_api_base() + if "ollama" in provider: + api_base = api_base.rstrip('/v1') + # Create response object for streaming if stream: - # For streaming with Ollama, we need to create a Response object first response = Response( id=FAKE_RESPONSES_ID, created_at=time.time(), model=self.model, object="response", output=[], - tool_choice="auto" if tool_choice is None or tool_choice == NOT_GIVEN else cast(Literal["auto", "required", "none"], tool_choice), + tool_choice="auto" if tool_choice is None or tool_choice == NOT_GIVEN else + cast(Literal["auto", "required", "none"], tool_choice), top_p=model_settings.top_p, temperature=model_settings.temperature, tools=[], parallel_tool_calls=parallel_tool_calls or False, ) - - # Special handling for Qwen when streaming to ensure tool calls are properly handled - if is_qwen and os.getenv('CAI_STREAM', 'false').lower() == 'true': - # Make sure we have the right custom provider in streaming mode - stream_obj = await litellm.acompletion( - **ollama_kwargs, - api_base=get_ollama_api_base(), - custom_llm_provider="openai" # Use openai provider for best compatibility with Qwen - ) - else: - # Standard Ollama streaming - stream_obj = await litellm.acompletion( - **ollama_kwargs, - api_base=get_ollama_api_base().rstrip('/v1'), - custom_llm_provider="ollama" - ) + # Get streaming response + stream_obj = await litellm.acompletion( + **ollama_kwargs, + api_base=api_base, + custom_llm_provider=provider, + ) return response, stream_obj else: - # Non-streaming mode - # For Qwen models, use the openai provider when tools are present - if is_qwen and "tools" in ollama_kwargs: - ret = litellm.completion( - **ollama_kwargs, - api_base=get_ollama_api_base(), - custom_llm_provider="openai" # Use openai provider for better tool handling - ) - else: - # Standard Ollama completion - ret = litellm.completion( - **ollama_kwargs, - api_base=get_ollama_api_base().rstrip('/v1'), - custom_llm_provider="ollama" - ) - return ret + + + # Get completion response + return litellm.completion( + **ollama_kwargs, + api_base=api_base, + custom_llm_provider=provider, + ) def _get_client(self) -> AsyncOpenAI: if self._client is None: @@ -1494,6 +1814,33 @@ class OpenAIChatCompletionsModel(Model): 'arguments': function_call.get('arguments', '') } }] + + # Handle special Ollama generic_linux_command format + if isinstance(delta, dict) and 'content' in delta: + content = delta['content'] + # Try to detect if the content is a JSON string with function call format + try: + if isinstance(content, str) and '{' in content and '}' in content: + # Try to extract JSON from the content (it might be embedded in text) + json_start = content.find('{') + json_end = content.rfind('}') + 1 + if json_start >= 0 and json_end > json_start: + json_str = content[json_start:json_end] + parsed = json.loads(json_str) + if 'name' in parsed and 'arguments' in parsed: + # This looks like a function call in JSON format + return [{ + 'index': 0, + 'id': f"call_{time.time_ns()}", # Generate a unique ID + 'type': 'function', + 'function': { + 'name': parsed['name'], + 'arguments': json.dumps(parsed['arguments']) if isinstance(parsed['arguments'], dict) else parsed['arguments'] + } + }] + except Exception: + # If JSON parsing fails, just continue with normal processing + pass # Anthropic-style tool_use format if hasattr(delta, 'tool_use') and delta.tool_use: diff --git a/src/cai/sdk/agents/run.py b/src/cai/sdk/agents/run.py index 84bb664b..a0937d28 100644 --- a/src/cai/sdk/agents/run.py +++ b/src/cai/sdk/agents/run.py @@ -4,6 +4,7 @@ import asyncio import copy from dataclasses import dataclass, field from typing import Any, cast +import os from openai.types.responses import ResponseCompletedEvent @@ -778,6 +779,23 @@ class Runner: run_config: RunConfig, tool_use_tracker: AgentToolUseTracker, ) -> SingleStepResult: + # Log the raw model response output, focusing on tool calls + logger.debug(f"[_get_single_step_result_from_response] Raw new_response.id: {new_response.referenceable_id}") + if new_response.output: + for i, item in enumerate(new_response.output): + item_type = type(item).__name__ + logger.debug(f"[_get_single_step_result_from_response] Raw output item [{i}] type: {item_type}") + if hasattr(item, "name") and hasattr(item, "arguments"): # For ResponseFunctionToolCall + logger.debug( + f"[_get_single_step_result_from_response] Raw ResponseFunctionToolCall item [{i}]: " + f"Name='{getattr(item, 'name', 'N/A')}', " + f"Args='{getattr(item, 'arguments', 'N/A')}', " + f"CallID='{getattr(item, 'call_id', 'N/A')}'" + ) + elif hasattr(item, "text"): # For ResponseOutputText + logger.debug(f"[_get_single_step_result_from_response] Raw ResponseOutputText item [{i}]: Text='{getattr(item, 'text', '')[:100]}...'") + + processed_response = RunImpl.process_model_response( agent=agent, all_tools=all_tools, @@ -786,6 +804,69 @@ class Runner: handoffs=handoffs, ) + # Log the processed response, focusing on tools_used + logger.debug(f"[_get_single_step_result_from_response] Processed response type: {type(processed_response).__name__}") + if hasattr(processed_response, 'is_final_output'): + logger.debug(f"[_get_single_step_result_from_response] Processed response: is_final_output={processed_response.is_final_output}") + else: + logger.debug(f"[_get_single_step_result_from_response] Processed response does not have is_final_output attribute") + + if hasattr(processed_response, 'final_output_from_llm') and processed_response.final_output_from_llm is not None: + logger.debug(f"[_get_single_step_result_from_response] Processed response: final_output_from_llm='{str(processed_response.final_output_from_llm)[:100]}...'") + + # Log tools used with robust type checking + if hasattr(processed_response, 'tools_used') and processed_response.tools_used: + # Log summarizing number of tools used first + logger.debug(f"[_get_single_step_result_from_response] Found {len(processed_response.tools_used)} tools used") + + # Add spacing between blocks in terminal output + if os.environ.get('CAI_STREAM', 'false').lower() == 'true': + print("") # Visual separator only when streaming is enabled + + for i, tool_call in enumerate(processed_response.tools_used): + try: + # Safely extract tool name with multiple fallbacks + tool_name = "Unknown" + try: + if hasattr(tool_call, 'tool'): + if isinstance(tool_call.tool, str): + tool_name = tool_call.tool + elif hasattr(tool_call.tool, 'name'): + tool_name = tool_call.tool.name + else: + tool_name = str(tool_call.tool) + except Exception: + pass + + # Safely extract call_id + call_id = "Unknown" + try: + if hasattr(tool_call, 'call_id'): + call_id = str(tool_call.call_id) + except Exception: + pass + + # Safely extract parsed_args + parsed_args = "Unknown" + try: + if hasattr(tool_call, 'parsed_args'): + parsed_args = str(tool_call.parsed_args) + except Exception: + pass + + logger.debug( + f"[_get_single_step_result_from_response] Processed tool_call [{i}]: " + f"Name='{tool_name}', " + f"CallID='{call_id}', " + f"ParsedArgs='{parsed_args}'" + ) + except Exception as e: + # Last resort fallback for any unexpected structure + logger.debug(f"[_get_single_step_result_from_response] Could not process tool_call [{i}]: {str(e)}") + else: + logger.debug("[_get_single_step_result_from_response] Processed response: No tools_used.") + + tool_use_tracker.add_tool_use(agent, processed_response.tools_used) return await RunImpl.execute_tools_and_side_effects( diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index e1e58778..95b0f223 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -222,18 +222,41 @@ def terminate_session(session_id): 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 - output = ctf.get_shell(command, timeout=timeout) - # exploit_logger.log_ok() - - if stdout: - print("\033[32m" + output + "\033[0m") - return output # output if output else result.stder + # Check the type of ctf object to handle various formats properly + if ctf is None: + return f"Error: CTF environment is None" + + # Handle string CTF values (like "CTF_ENV") + if isinstance(ctf, str): + # If in local mode, fallback to running local command + return _run_local(command, stdout, timeout, stream, call_id) + + # Handle dict format that might come from Ollama/Qwen models + if isinstance(ctf, dict): + # Check if this dict has a get_shell key or function + if callable(getattr(ctf, 'get_shell', None)): + # This is a proper CTF object + output = ctf.get_shell(command, timeout=timeout) + if stdout: + print("\033[32m" + output + "\033[0m") + return output + else: + # This is a dict but not a proper CTF object, fallback to local + return _run_local(command, stdout, timeout, stream, call_id) + + # Original code path for proper CTF objects + if hasattr(ctf, 'get_shell') and callable(ctf.get_shell): + output = ctf.get_shell(command, timeout=timeout) + if stdout: + print("\033[32m" + output + "\033[0m") + return output + + # Fallback for any other case + return _run_local(command, stdout, timeout, stream, call_id) except Exception as e: # pylint: disable=broad-except print(color(f"Error executing CTF command: {e}", fg="red")) - # exploit_logger.log_error(str(e)) - return f"Error executing CTF command: {str(e)}" + # Fallback to local execution + return _run_local(command, stdout, timeout, stream, call_id) def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None): @@ -328,7 +351,7 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): header.append(")", style="yellow") tool_time = 0 start_time = time.time() - total_time = time.time() - START_TIME + total_time = time.time() - START_TIME if START_TIME else 0 timing_info = [] if total_time: timing_info.append(f"Total: {format_time(total_time)}") @@ -339,15 +362,24 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): content = Text() + # Get console width and calculate appropriate panel width + console_width = console.width if hasattr(console, 'width') else 100 + panel_width = min(int(console_width * 0.95), console_width - 4) # Use 95% of width or leave 4 chars margin + + # Create the panel with a specific width to avoid overflow panel = Panel( Text.assemble(header, "\n\n", content), title="[bold green]Tool Execution[/bold green]", subtitle="[bold green]Live Output[/bold green]", border_style="green", padding=(1, 2), - box=ROUNDED + box=ROUNDED, + width=panel_width ) + # Print clear separator before panel to avoid overlap with previous content + console.print("\n\n") + # Start Live display with Live(panel, console=console, refresh_per_second=4) as live: # Stream stdout in real-time @@ -364,7 +396,7 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): # Update tool_time and header with new timing info tool_time = time.time() - start_time - total_time = time.time() - START_TIME + total_time = time.time() - START_TIME if START_TIME else 0 # Remove any previous timing info from header (rebuild header) timing_info = [] if total_time: @@ -386,7 +418,8 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): subtitle="[bold green]Live Output[/bold green]", border_style="green", padding=(1, 2), - box=ROUNDED + box=ROUNDED, + width=panel_width ) live.update(panel) # Check if process is done @@ -405,7 +438,8 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): subtitle="[bold green]Live Output[/bold green]", border_style="green", padding=(1, 2), - box=ROUNDED + box=ROUNDED, + width=panel_width ) live.update(panel) @@ -415,12 +449,16 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): title="[bold green]Tool Execution[/bold green]", border_style="green", padding=(1, 2), - box=ROUNDED + box=ROUNDED, + width=panel_width ) live.update(panel) # Wait a moment for the panel to be displayed properly time.sleep(0.5) + + # Print clear separator after panel to avoid overlap with next content + console.print("\n\n") else: # Fallback to simpler streaming with cli_print_tool_output # Parse command into command and args (same as rich mode) @@ -436,6 +474,9 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): tool_args["args"] = args # Note: Omitted empty values and async_mode=False as it's default + # Print separator to avoid overlap with previous content + print("\n") + # Initial notification - just once cli_print_tool_output(tool_name, tool_args, "Command started...", call_id=call_id) @@ -473,6 +514,9 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): final_output += f"\nCommand exited with code {return_code}" cli_print_tool_output(tool_name, tool_args, final_output, call_id=call_id) + + # Print separator to avoid overlap with next content + print("\n") # Return the full output return ''.join(output_buffer) @@ -547,7 +591,22 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg 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, stream, call_id) - return _run_local(command, stdout, timeout, stream, call_id, tool_name) + # Determine whether to use CTF or local execution + use_ctf = False + if ctf is not None: + # Check if ctf is a proper object that can handle shell commands + if (hasattr(ctf, 'get_shell') and callable(ctf.get_shell)) or isinstance(ctf, dict) or isinstance(ctf, str): + use_ctf = True + + # Run the command using the appropriate method + if use_ctf: + try: + # Try with CTF first + return _run_ctf(ctf, command, stdout, timeout, stream, call_id) + except Exception as e: # pylint: disable=broad-except + # Fallback to local if CTF fails + print(color(f"CTF execution failed, falling back to local: {str(e)}", fg="yellow")) + return _run_local(command, stdout, timeout, stream, call_id, tool_name) + else: + # Use local execution + return _run_local(command, stdout, timeout, stream, call_id, tool_name) From 42aed4ea7055880cd9e4c9164dc5ddab61dc714e Mon Sep 17 00:00:00 2001 From: luijait Date: Tue, 6 May 2025 16:23:10 +0200 Subject: [PATCH 21/29] Qwen executes tool calls --- src/cai/sdk/agents/run.py | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/src/cai/sdk/agents/run.py b/src/cai/sdk/agents/run.py index a0937d28..3c45b72f 100644 --- a/src/cai/sdk/agents/run.py +++ b/src/cai/sdk/agents/run.py @@ -853,16 +853,6 @@ class Runner: parsed_args = str(tool_call.parsed_args) except Exception: pass - - logger.debug( - f"[_get_single_step_result_from_response] Processed tool_call [{i}]: " - f"Name='{tool_name}', " - f"CallID='{call_id}', " - f"ParsedArgs='{parsed_args}'" - ) - except Exception as e: - # Last resort fallback for any unexpected structure - logger.debug(f"[_get_single_step_result_from_response] Could not process tool_call [{i}]: {str(e)}") else: logger.debug("[_get_single_step_result_from_response] Processed response: No tools_used.") From 2c6fa34cd6dba2aaf2c0809dc7e1163fd968746e Mon Sep 17 00:00:00 2001 From: luijait Date: Tue, 6 May 2025 16:47:39 +0200 Subject: [PATCH 22/29] Qwen executes tool calls --- src/cai/sdk/agents/run.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/cai/sdk/agents/run.py b/src/cai/sdk/agents/run.py index 3c45b72f..30d82d2f 100644 --- a/src/cai/sdk/agents/run.py +++ b/src/cai/sdk/agents/run.py @@ -853,9 +853,8 @@ class Runner: parsed_args = str(tool_call.parsed_args) except Exception: pass - else: - logger.debug("[_get_single_step_result_from_response] Processed response: No tools_used.") - + except Exception: + pass tool_use_tracker.add_tool_use(agent, processed_response.tools_used) From a943938e493ad8b1fc0f94224b779d58d5ef2da7 Mon Sep 17 00:00:00 2001 From: luijait Date: Tue, 6 May 2025 17:11:18 +0200 Subject: [PATCH 23/29] generalice tool calls, remove ctf from arguments --- .../agents/models/openai_chatcompletions.py | 251 ++---------------- 1 file changed, 24 insertions(+), 227 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index c5138d14..21c5ef81 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -809,154 +809,48 @@ class OpenAIChatCompletionsModel(Model): # Try to extract a JSON object from the content json_start = ollama_full_content.find('{') json_end = ollama_full_content.rfind('}') + 1 - - logger.debug(f"Ollama content length: {len(ollama_full_content)}, JSON start: {json_start}, JSON end: {json_end}") - + if json_start >= 0 and json_end > json_start: - json_str = ollama_full_content[json_start:json_end] - logger.debug(f"Extracted potential JSON: {json_str[:100]}...") - - # Special check for generic_linux_command format - if "generic_linux_command" in json_str and "arguments" in json_str: - logger.debug("Detected generic_linux_command pattern - using special handling") - - # Try regex pattern matching to handle potentially malformed JSON - cmd_pattern = re.search(r'"command"\s*:\s*"([^"]+)"', json_str) - args_pattern = re.search(r'"args"\s*:\s*"([^"]*)"', json_str) - ctf_pattern = re.search(r'"ctf"\s*:\s*"([^"]*)"', json_str) - - if cmd_pattern: - # Create a tool call specifically for generic_linux_command - command = cmd_pattern.group(1) - args = args_pattern.group(1) if args_pattern else "" - ctf = ctf_pattern.group(1) if ctf_pattern else "" - - tool_call_id = f"call_{hashlib.md5(('generic_linux_command' + str(time.time())).encode()).hexdigest()[:8]}" - - # Create a proper arguments object - command_args = { - "command": command, - "args": args, - "ctf": ctf, - "async_mode": False, - "session_id": "" - } - - arguments_str = json.dumps(command_args) - logger.debug(f"Created generic_linux_command with args: {arguments_str}") - - # Add it to our function_calls state - state.function_calls[0] = ResponseFunctionToolCall( - id=FAKE_RESPONSES_ID, - arguments=arguments_str, - name="generic_linux_command", - type="function_call", - call_id=tool_call_id, - ) - - # Display the tool call in CLI - from cai.util import cli_print_agent_messages - try: - # Create a message-like object to display the function call - tool_msg = type('ToolCallWrapper', (), { - 'content': None, - 'tool_calls': [ - type('ToolCallDetail', (), { - 'function': type('FunctionDetail', (), { - 'name': "generic_linux_command", - 'arguments': arguments_str - }), - 'id': tool_call_id, - 'type': 'function' - }) - ] - }) - - # Print the tool call using the CLI utility - cli_print_agent_messages( - agent_name=getattr(self, 'agent_name', 'Agent'), - message=tool_msg, - counter=getattr(self, 'interaction_counter', 0), - model=str(self.model), - debug=False, - interaction_input_tokens=estimated_input_tokens, - interaction_output_tokens=estimated_output_tokens, - interaction_reasoning_tokens=0, - total_input_tokens=getattr(self, 'total_input_tokens', 0) + estimated_input_tokens, - total_output_tokens=getattr(self, 'total_output_tokens', 0) + estimated_output_tokens, - total_reasoning_tokens=0, - interaction_cost=None, - total_cost=None - ) - except Exception as e: - logger.error(f"Error displaying tool call in CLI: {e}") - - # Add to message history - tool_call_msg = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": tool_call_id, - "type": "function", - "function": { - "name": "generic_linux_command", - "arguments": arguments_str - } - } - ] - } - - streamed_tool_calls.append(tool_call_msg) - add_to_message_history(tool_call_msg) - logger.debug(f"Added generic_linux_command with args: {arguments_str}") - - - # Standard JSON parsing for other function calls + json_str = ollama_full_content[json_start:json_end] + # Try to parse the JSON parsed = json.loads(json_str) # Check if it looks like a function call - if ('name' in parsed and 'arguments' in parsed and - isinstance(parsed['arguments'], dict)): - + if ('name' in parsed and 'arguments' in parsed): logger.debug(f"Found valid function call in Ollama output: {json_str}") - # Create a tool call from the JSON object + # Create a tool call ID tool_call_id = f"call_{hashlib.md5((parsed['name'] + str(time.time())).encode()).hexdigest()[:8]}" - # Add it to our function_calls state - state.function_calls[0] = ResponseFunctionToolCall( - id=FAKE_RESPONSES_ID, - arguments=json.dumps(parsed['arguments']), - name=parsed['name'], - type="function_call", - call_id=tool_call_id, - ) - # Ensure arguments is a valid JSON string arguments_str = "" if isinstance(parsed['arguments'], dict): + # Remove 'ctf' field if it exists + if 'ctf' in parsed['arguments']: + del parsed['arguments']['ctf'] arguments_str = json.dumps(parsed['arguments']) elif isinstance(parsed['arguments'], str): # If it's already a string, check if it's valid JSON try: - # Try parsing to validate - json.loads(parsed['arguments']) - arguments_str = parsed['arguments'] + # Try parsing to validate and remove 'ctf' if present + args_dict = json.loads(parsed['arguments']) + if isinstance(args_dict, dict) and 'ctf' in args_dict: + del args_dict['ctf'] + arguments_str = json.dumps(args_dict) except: # If not valid JSON, encode it as a JSON string arguments_str = json.dumps(parsed['arguments']) else: # For any other type, convert to string and then JSON - arguments_str = json.dumps(str(parsed['arguments'])) - - logger.debug(f"Final arguments string: {arguments_str}") - - # Update with the properly formatted arguments - state.function_calls[0].arguments = arguments_str - - # Log the tool call for development purposes - logger.debug(f"Adding tool call: {parsed['name']} with arguments: {arguments_str}") + arguments_str = json.dumps(str(parsed['arguments'])) + # Add it to our function_calls state + state.function_calls[0] = ResponseFunctionToolCall( + id=FAKE_RESPONSES_ID, + arguments=arguments_str, + name=parsed['name'], + type="function_call", + call_id=tool_call_id, + ) # Display the tool call in CLI from cai.util import cli_print_agent_messages @@ -1015,106 +909,9 @@ class OpenAIChatCompletionsModel(Model): streamed_tool_calls.append(tool_call_msg) add_to_message_history(tool_call_msg) - logger.debug(f"Added function call: {parsed['name']} with args: {json.dumps(parsed['arguments'])}") + logger.debug(f"Added function call: {parsed['name']} with args: {arguments_str}") except Exception as e: - logger.debug(f"Failed to parse potential Ollama function call: {e}") - - # Even if JSON parsing fails, try to extract generic_linux_command with regex - if "generic_linux_command" in ollama_full_content: - logger.debug("JSON parsing failed but detected generic_linux_command - trying regex fallback") - try: - # Use regex to extract command components even from malformed output - cmd_pattern = re.search(r'"command"\s*:\s*"([^"]+)"', ollama_full_content) - args_pattern = re.search(r'"args"\s*:\s*"([^"]*)"', ollama_full_content) - ctf_pattern = re.search(r'"ctf"\s*:\s*"([^"]*)"', ollama_full_content) - - if cmd_pattern: - # Create a tool call specifically for generic_linux_command - command = cmd_pattern.group(1) - args = args_pattern.group(1) if args_pattern else "" - ctf = ctf_pattern.group(1) if ctf_pattern else "" - - tool_call_id = f"call_{hashlib.md5(('generic_linux_command' + str(time.time())).encode()).hexdigest()[:8]}" - - # Create a proper arguments object - command_args = { - "command": command, - "args": args, - "ctf": ctf, - "async_mode": False, - "session_id": "" - } - - arguments_str = json.dumps(command_args) - logger.debug(f"Fallback created generic_linux_command with args: {arguments_str}") - - # Add it to our function_calls state - state.function_calls[0] = ResponseFunctionToolCall( - id=FAKE_RESPONSES_ID, - arguments=arguments_str, - name="generic_linux_command", - type="function_call", - call_id=tool_call_id, - ) - - # Display the tool call in CLI - from cai.util import cli_print_agent_messages - try: - # Create a message-like object for display - tool_msg = type('ToolCallWrapper', (), { - 'content': None, - 'tool_calls': [ - type('ToolCallDetail', (), { - 'function': type('FunctionDetail', (), { - 'name': "generic_linux_command", - 'arguments': arguments_str - }), - 'id': tool_call_id, - 'type': 'function' - }) - ] - }) - - cli_print_agent_messages( - agent_name=getattr(self, 'agent_name', 'Agent'), - message=tool_msg, - counter=getattr(self, 'interaction_counter', 0), - model=str(self.model), - debug=False, - interaction_input_tokens=estimated_input_tokens, - interaction_output_tokens=estimated_output_tokens, - interaction_reasoning_tokens=0, - total_input_tokens=getattr(self, 'total_input_tokens', 0) + estimated_input_tokens, - total_output_tokens=getattr(self, 'total_output_tokens', 0) + estimated_output_tokens, - total_reasoning_tokens=0, - interaction_cost=None, - total_cost=None - ) - except Exception as e: - logger.error(f"Error displaying fallback tool call in CLI: {e}") - - # Add to message history - tool_call_msg = { - "role": "assistant", - "content": None, - "tool_calls": [ - { - "id": tool_call_id, - "type": "function", - "function": { - "name": "generic_linux_command", - "arguments": arguments_str - } - } - ] - } - - streamed_tool_calls.append(tool_call_msg) - add_to_message_history(tool_call_msg) - - logger.debug(f"Added fallback generic_linux_command") - except Exception as regex_err: - logger.error(f"Regex fallback also failed: {regex_err}") + pass function_call_starting_index = 0 if state.text_content_index_and_output: From c378d7375b0e09aa62ba3d1a1c17970ccb430fb5 Mon Sep 17 00:00:00 2001 From: luijait Date: Tue, 6 May 2025 17:14:01 +0200 Subject: [PATCH 24/29] rebase common.py --- src/cai/tools/common.py | 601 +++++++++++++++++++++++++++++----------- 1 file changed, 439 insertions(+), 162 deletions(-) diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index 95b0f223..e475d5f4 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -24,14 +24,70 @@ except ImportError: # Global dictionary to store active sessions ACTIVE_SESSIONS = {} +def _get_workspace_dir() -> str: + """Determines the target workspace directory based on env vars for host.""" + base_dir_env = os.getenv("CAI_WORKSPACE_DIR") + workspace_name = os.getenv("CAI_WORKSPACE") + + # Determine the base directory + if base_dir_env: + base_dir = os.path.abspath(base_dir_env) + else: # Default base directory is 'workspaces' + if workspace_name: + base_dir = os.path.join(os.getcwd(), "workspaces") + else: # If no workspace name is set, the workspace IS the CWD. + return os.getcwd() + + # If a workspace name is provided, append it to the base directory + if workspace_name: + if not all(c.isalnum() or c in ['_', '-'] for c in workspace_name): + print(color(f"Invalid CAI_WORKSPACE name '{workspace_name}'. " + f"Using directory '{base_dir}' instead.", fg="yellow")) + target_dir = base_dir + else: + target_dir = os.path.join(base_dir, workspace_name) + else: + target_dir = base_dir + + # Ensure the final target directory exists on the host + try: + abs_target_dir = os.path.abspath(target_dir) + os.makedirs(abs_target_dir, exist_ok=True) + return abs_target_dir + except OSError as e: + print(color(f"Error creating/accessing host workspace directory '{abs_target_dir}': {e}", + fg="red")) + print(color(f"Falling back to current directory: {os.getcwd()}", fg="yellow")) + return os.getcwd() + +def _get_container_workspace_path() -> str: + """Determines the target workspace path inside the container.""" + workspace_name = os.getenv("CAI_WORKSPACE") + if workspace_name: + if not all(c.isalnum() or c in ['_', '-'] for c in workspace_name): + print(color(f"Invalid CAI_WORKSPACE name '{workspace_name}' for container. " + f"Using '/workspace'.", fg="yellow")) + return "/" + # Standard path inside CAI containers + return f"/workspace/workspaces/{workspace_name}" + else: + return "/" class ShellSession: # pylint: disable=too-many-instance-attributes """Class to manage interactive shell sessions""" - def __init__(self, command, session_id=None, ctf=None): + def __init__(self, command, session_id=None, ctf=None, workspace_dir=None, container_id=None): # noqa E501 self.session_id = session_id or str(uuid.uuid4())[:8] - self.command = command + self.command = command self.ctf = ctf + self.container_id = container_id + # Determine workspace based on context (container, ctf or local host) + if self.container_id: + self.workspace_dir = _get_container_workspace_path() + elif self.ctf: + self.workspace_dir = workspace_dir or _get_workspace_dir() + else: + self.workspace_dir = _get_workspace_dir() self.process = None self.master = None self.slave = None @@ -40,9 +96,40 @@ class ShellSession: # pylint: disable=too-many-instance-attributes self.last_activity = time.time() def start(self): - """Start the shell session""" + """Start the shell session in the appropriate environment.""" + start_message_cmd = self.command + + # --- Start in Container --- + if self.container_id: + try: + self.master, self.slave = pty.openpty() + docker_cmd_list = [ + "docker", "exec", "-i", + "-w", self.workspace_dir, + self.container_id, + "sh", "-c", # Use shell to handle complex commands if needed + self.command # The actual command to run + ] + self.process = subprocess.Popen( + docker_cmd_list, + stdin=self.slave, + stdout=self.slave, + stderr=self.slave, + preexec_fn=os.setsid, + universal_newlines=True + ) + self.is_running = True + self.output_buffer.append( + f"[Session {self.session_id}] Started in container {self.container_id[:12]}: " + f"{start_message_cmd} in {self.workspace_dir}") + threading.Thread(target=self._read_output, daemon=True).start() + except Exception as e: + self.output_buffer.append(f"Error starting container session: {str(e)}") + self.is_running = False + return + + # --- Start in CTF --- if self.ctf: - # For CTF environments self.is_running = True self.output_buffer.append( f"[Session { @@ -52,81 +139,100 @@ class ShellSession: # pylint: disable=too-many-instance-attributes output = self.ctf.get_shell(self.command) self.output_buffer.append(output) except Exception as e: # pylint: disable=broad-except - self.output_buffer.append(f"Error: {str(e)}") - self.is_running = False + self.output_buffer.append(f"Error executing CTF command: {str(e)}") + self.is_running = False return - # For local environment + # --- Start Locally (Host) --- try: - # Create a pseudo-terminal self.master, self.slave = pty.openpty() - - # Start the process self.process = subprocess.Popen( # pylint: disable=subprocess-popen-preexec-fn, consider-using-with # noqa: E501 - self.command, - shell=True, # nosec B602 + self.command, + shell=True, # nosec B602 stdin=self.slave, stdout=self.slave, stderr=self.slave, - preexec_fn=os.setsid, # Create a new process group + cwd=self.workspace_dir, + preexec_fn=os.setsid, universal_newlines=True ) - self.is_running = True self.output_buffer.append( f"[Session { self.session_id}] Started: { self.command}") - # Start a thread to read output threading.Thread(target=self._read_output, daemon=True).start() except Exception as e: # pylint: disable=broad-except - self.output_buffer.append(f"Error starting session: {str(e)}") + self.output_buffer.append(f"Error starting local session: {str(e)}") self.is_running = False def _read_output(self): """Read output from the process""" try: - while self.is_running: + while self.is_running and self.master is not None: try: + # Check if process has exited before reading + if self.process and self.process.poll() is not None: + self.is_running = False + break + # Read the output output = os.read(self.master, 1024).decode() if output: self.output_buffer.append(output) self.last_activity = time.time() - except OSError: - # No data available or terminal closed - time.sleep(0.1) - if not self.is_process_running(): + else: self.is_running = False break - except Exception as e: # pylint: disable=broad-except - self.output_buffer.append(f"Error reading output: {str(e)}") + except Exception as e: + self.output_buffer.append(f"Error reading output buffer: {str(read_err)}") + self.is_running = False + break + # Add a small sleep to prevent busy-waiting if no output + if is_process_running(self): + time.sleep(0.05) + except Exception as e: + self.output_buffer.append(f"Error in read_output loop: {str(e)}") self.is_running = False + def is_process_running(self): """Check if the process is still running""" + # For CTF or container + if self.container_id or self.ctf: + return self.is_running + # For local host if not self.process: return False return self.process.poll() is None def send_input(self, input_data): - """Send input to the process""" - if not self.is_running: - return "Session is not running" + """Send input to the process (local or container)""" + if not self.is_running: # For CTF or container + if self.process and self.process.poll() is None: + self.is_running = True + else: # For local host + return "Session is not running" try: + # --- Send to CTF --- if self.ctf: - # For CTF environments output = self.ctf.get_shell(input_data) self.output_buffer.append(output) return "Input sent to CTF session" - # For local environment - input_data = input_data.rstrip() + "\n" - os.write(self.master, input_data.encode()) - self.last_activity = time.time() - return "Input sent to session" + # --- Send to Local or Container PTY --- + if self.master is not None: + input_data_bytes = (input_data.rstrip() + "\n").encode() + bytes_written = os.write(self.master, input_data_bytes) + if bytes_written != len(input_data_bytes): + self.output_buffer.append(f"[Session {self.session_id}] Warning: Partial input write.") + self.last_activity = time.time() + return "Input sent to session" + else: + return "Session PTY not available for input" except Exception as e: # pylint: disable=broad-except + self.output_buffer.append(f"Error sending input: {str(e)}") return f"Error sending input: {str(e)}" def get_output(self, clear=True): @@ -138,8 +244,12 @@ class ShellSession: # pylint: disable=too-many-instance-attributes def terminate(self): """Terminate the session""" + session_id_short = self.session_id[:8] if not self.is_running: - return "Session already terminated" + if self.process and self.process.poll() is None: + pass # Process is running, proceed with termination + else: + return f"Session {session_id_short} already terminated or finished." try: self.is_running = False @@ -147,28 +257,64 @@ class ShellSession: # pylint: disable=too-many-instance-attributes if self.process: # Try to terminate the process group try: - os.killpg(os.getpgid(self.process.pid), signal.SIGTERM) - except BaseException: # pylint: disable=bare-except,broad-except # noqa: E501 - # If that fails, try to terminate just the process - self.process.terminate() + os.killpg(os.getpgid(self.process.pid), signal.SIGTERM) + except ProcessLookupError: + pass # Process already gone + except subprocess.TimeoutExpired: + print(color(f"Session {session_id_short} did not terminate gracefully, sending SIGKILL...", fg="yellow")) # noqa E501 + try: + if pgid: + os.killpg(pgid, signal.SIGKILL) # Force kill + else: + self.process.kill() + except ProcessLookupError: + pass # Already gone + except Exception as kill_err: + termination_message = f" (Error during SIGKILL: {kill_err})" + except Exception as term_err: # Catch other errors during SIGTERM + termination_message = f" (Error during SIGTERM: {term_err})" + try: + self.process.kill() + except Exception: pass # Ignore nested errors - # Clean up resources - if self.master: - os.close(self.master) - if self.slave: - os.close(self.slave) - return f"Session {self.session_id} terminated" + # Final check + if self.process.poll() is None: + print(color(f"Session {session_id_short} process {self.process.pid} may still be running after termination attempts.", fg="red")) # noqa E501 + termination_message += " (Warning: Process may still be running)" + + + # Clean up PTY resources if they exist + if self.master: + try: os.close(self.master) + except OSError: pass + self.master = None + if self.slave: + try: os.close(self.slave) + except OSError: pass + self.slave = None + + return termination_message or f"Session {self.session_id} terminated" except Exception as e: # pylint: disable=broad-except - return f"Error terminating session: {str(e)}" + return f"Error terminating session {session_id_short}: {str(e)}" -def create_shell_session(command, ctf=None): - """Create a new shell session""" - session = ShellSession(command, ctf=ctf) +def create_shell_session(command, ctf=None, container_id=None, **kwargs): + """Create a new shell session in the correct workspace/environment.""" + if container_id: + session = ShellSession(command, ctf=ctf, container_id=container_id) + else: + workspace_dir = _get_workspace_dir() + session = ShellSession(command, ctf=ctf, workspace_dir=workspace_dir) + session.start() - ACTIVE_SESSIONS[session.session_id] = session - return session.session_id + if session.is_running or (ctf and not session.is_running): + ACTIVE_SESSIONS[session.session_id] = session + return session.session_id + else: + error_msg = session.get_output(clear=True) + print(color(f"Failed to start session: {error_msg}", fg="red")) + return f"Failed to start session: {error_msg}" def list_shell_sessions(): @@ -212,70 +358,100 @@ def get_session_output(session_id, clear=True): def terminate_session(session_id): """Terminate a specific session""" if session_id not in ACTIVE_SESSIONS: - return f"Session {session_id} not found" + return f"Session {session_id} not found or already terminated." session = ACTIVE_SESSIONS[session_id] result = session.terminate() - del ACTIVE_SESSIONS[session_id] + if session_id in ACTIVE_SESSIONS: + del ACTIVE_SESSIONS[session_id] return result -def _run_ctf(ctf, command, stdout=False, timeout=100, stream=False, call_id=None): +def _run_ctf(ctf, command, stdout=False, timeout=100, workspace_dir=None): + """Runs command in CTF env, changing to workspace_dir first.""" + target_dir = workspace_dir or _get_workspace_dir() + full_command = f"cd '{target_dir}' && {command}" + original_cmd_for_msg = command # For logging + context_msg = f"(ctf:{target_dir})" try: - # Check the type of ctf object to handle various formats properly - if ctf is None: - return f"Error: CTF environment is None" - - # Handle string CTF values (like "CTF_ENV") - if isinstance(ctf, str): - # If in local mode, fallback to running local command - return _run_local(command, stdout, timeout, stream, call_id) - - # Handle dict format that might come from Ollama/Qwen models - if isinstance(ctf, dict): - # Check if this dict has a get_shell key or function - if callable(getattr(ctf, 'get_shell', None)): - # This is a proper CTF object - output = ctf.get_shell(command, timeout=timeout) - if stdout: - print("\033[32m" + output + "\033[0m") - return output - else: - # This is a dict but not a proper CTF object, fallback to local - return _run_local(command, stdout, timeout, stream, call_id) - - # Original code path for proper CTF objects - if hasattr(ctf, 'get_shell') and callable(ctf.get_shell): - output = ctf.get_shell(command, timeout=timeout) - if stdout: - print("\033[32m" + output + "\033[0m") - return output - - # Fallback for any other case - return _run_local(command, stdout, timeout, stream, call_id) + output = ctf.get_shell(full_command, timeout=timeout) + if stdout: + print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501 + return output except Exception as e: # pylint: disable=broad-except - print(color(f"Error executing CTF command: {e}", fg="red")) - # Fallback to local execution - return _run_local(command, stdout, timeout, stream, call_id) + error_msg = f"Error executing CTF command '{original_cmd_for_msg}' in '{target_dir}': {e}" # noqa E501 + print(color(error_msg, fg="red")) + return error_msg + +def _run_ssh(command, stdout=False, timeout=100, workspace_dir=None): + """Runs command via SSH. Assumes SSH agent or passwordless setup unless sshpass is used externally.""" # noqa E501 + ssh_user = os.environ.get('SSH_USER') + ssh_host = os.environ.get('SSH_HOST') + ssh_pass = os.environ.get('SSH_PASS') + remote_command = command + original_cmd_for_msg = command + context_msg = f"({ssh_user}@{ssh_host})" + + # Construct base SSH command list + if ssh_pass: + ssh_cmd_list = ["sshpass", "-p", ssh_pass, "ssh", f"{ssh_user}@{ssh_host}"] # noqa E501 + else: + ssh_cmd_list = ["ssh", f"{ssh_user}@{ssh_host}"] + ssh_cmd_list.append(remote_command) + + try: + # Use subprocess.run with list of args for better security than shell=True + result = subprocess.run( + ssh_cmd_list, + capture_output=True, + text=True, + check=False, # Don't raise exception on non-zero exit code + timeout=timeout + ) + output = result.stdout if result.stdout else result.stderr + if stdout: + print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501 + # Return combined output, potentially including errors + return output.strip() + except subprocess.TimeoutExpired as e: + error_output = e.stdout if e.stdout else str(e) + timeout_msg = f"Timeout executing SSH command: {error_output}" + if stdout: + print(f"\033[33m{context_msg} $ {original_cmd_for_msg}\nTIMEOUT\n{error_output}\033[0m") # noqa E501 + return timeout_msg + except FileNotFoundError: + # Handle case where ssh or sshpass isn't installed + error_msg = f"'sshpass' or 'ssh' command not found. Ensure they are installed and in PATH." # noqa E501 + print(color(error_msg, fg="red")) + return error_msg + except Exception as e: # pylint: disable=broad-except + error_msg = f"Error executing SSH command '{original_cmd_for_msg}' on {ssh_host}: {e}" # noqa E501 + print(color(error_msg, fg="red")) + return error_msg -def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None): +def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None, workspace_dir=None): + """Runs command locally in the specified workspace_dir.""" # If streaming is enabled and we have a call_id if stream and call_id: - return _run_local_streamed(command, call_id, timeout, tool_name) + return _run_local_streamed(command, call_id, timeout, tool_name, workspace_dir) + target_dir = workspace_dir or _get_workspace_dir() + original_cmd_for_msg = command # For logging + context_msg = f"(local:{target_dir})" try: - # nosec B602 - shell=True is required for command chaining result = subprocess.run( command, shell=True, # nosec B602 capture_output=True, text=True, - check=False, - timeout=timeout) + check=False, + timeout=timeout, + cwd=target_dir + ) output = result.stdout if result.stdout else result.stderr if stdout: - print("\033[32m" + output + "\033[0m") + print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501 # Skip passing output to cli_print_tool_output when CAI_STREAM=true # This prevents duplicate output in streaming mode @@ -284,20 +460,21 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t # Optional: Add cli_print_tool_output call here if needed for non-streaming pass - return output + return output.strip() except subprocess.TimeoutExpired as e: - error_output = e.stdout.decode() if e.stdout else str(e) + error_output = e.stdout if e.stdout else str(e) if stdout: print("\033[32m" + error_output + "\033[0m") - return error_output + return error_output except Exception as e: # pylint: disable=broad-except - error_msg = f"Error executing local command: {e}" - print(color(error_msg, fg="red")) - return error_msg + error_msg = f"Error executing local command: {e}" + print(color(error_msg, fg="red")) + return error_msg -def _run_local_streamed(command, call_id, timeout=100, tool_name=None): - """Run a local command with streaming output to the Tool output panel""" +def _run_local_streamed(command, call_id, timeout=100, tool_name=None, workspace_dir=None): + """Run a local command with streaming output to the Tool output panel.""" + target_dir = workspace_dir or _get_workspace_dir() try: # Try to import Rich for nice display try: @@ -321,7 +498,8 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, - bufsize=1 + bufsize=1, + cwd=target_dir # Set CWD for local process ) # If tool_name is not provided, derive it from the command @@ -351,7 +529,7 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): header.append(")", style="yellow") tool_time = 0 start_time = time.time() - total_time = time.time() - START_TIME if START_TIME else 0 + total_time = time.time() - START_TIME timing_info = [] if total_time: timing_info.append(f"Total: {format_time(total_time)}") @@ -362,24 +540,15 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): content = Text() - # Get console width and calculate appropriate panel width - console_width = console.width if hasattr(console, 'width') else 100 - panel_width = min(int(console_width * 0.95), console_width - 4) # Use 95% of width or leave 4 chars margin - - # Create the panel with a specific width to avoid overflow panel = Panel( Text.assemble(header, "\n\n", content), title="[bold green]Tool Execution[/bold green]", subtitle="[bold green]Live Output[/bold green]", border_style="green", padding=(1, 2), - box=ROUNDED, - width=panel_width + box=ROUNDED ) - # Print clear separator before panel to avoid overlap with previous content - console.print("\n\n") - # Start Live display with Live(panel, console=console, refresh_per_second=4) as live: # Stream stdout in real-time @@ -396,7 +565,7 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): # Update tool_time and header with new timing info tool_time = time.time() - start_time - total_time = time.time() - START_TIME if START_TIME else 0 + total_time = time.time() - START_TIME # Remove any previous timing info from header (rebuild header) timing_info = [] if total_time: @@ -418,8 +587,7 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): subtitle="[bold green]Live Output[/bold green]", border_style="green", padding=(1, 2), - box=ROUNDED, - width=panel_width + box=ROUNDED ) live.update(panel) # Check if process is done @@ -438,8 +606,7 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): subtitle="[bold green]Live Output[/bold green]", border_style="green", padding=(1, 2), - box=ROUNDED, - width=panel_width + box=ROUNDED ) live.update(panel) @@ -449,16 +616,12 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): title="[bold green]Tool Execution[/bold green]", border_style="green", padding=(1, 2), - box=ROUNDED, - width=panel_width + box=ROUNDED ) live.update(panel) # Wait a moment for the panel to be displayed properly time.sleep(0.5) - - # Print clear separator after panel to avoid overlap with next content - console.print("\n\n") else: # Fallback to simpler streaming with cli_print_tool_output # Parse command into command and args (same as rich mode) @@ -474,9 +637,6 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): tool_args["args"] = args # Note: Omitted empty values and async_mode=False as it's default - # Print separator to avoid overlap with previous content - print("\n") - # Initial notification - just once cli_print_tool_output(tool_name, tool_args, "Command started...", call_id=call_id) @@ -514,9 +674,6 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None): final_output += f"\nCommand exited with code {return_code}" cli_print_tool_output(tool_name, tool_args, final_output, call_id=call_id) - - # Print separator to avoid overlap with next content - print("\n") # Return the full output return ''.join(output_buffer) @@ -549,7 +706,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg async_mode=False, session_id=None, timeout=100, stream=False, call_id=None, tool_name=None): """ - Run command either in CTF container or on the local attacker machine + Run command in the appropriate environment (Docker, CTF, SSH, Local) + and workspace. Args: command: The command to execute @@ -564,49 +722,168 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg If None, the tool name will be derived from the command. Returns: - str: Command output, status message, or session ID + str: Command output, status message, or session ID. """ # If session_id is provided, send command to that session if session_id: if session_id not in ACTIVE_SESSIONS: return f"Session {session_id} not found" - - result = send_to_session(session_id, command) + session = ACTIVE_SESSIONS[session_id] + result = session.send_input(command) # Send the raw command string if stdout: output = get_session_output(session_id, clear=False) - print("\033[32m" + output + "\033[0m") - return result - - # If async_mode, create a new session - if async_mode: - session_id = create_shell_session(command, ctf) - if stdout: - # Wait a moment for initial output - 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." + env_type = "Local" + if session.container_id: + env_type = f"Container({session.container_id[:12]})" + elif session.ctf: + env_type = "CTF" + print(f"\033[32m(Session {session_id} in {env_type}:{session.workspace_dir}) >> {command}\n{output}\033[0m") # noqa E501 + return result # Return the result of sending input ("Input sent..." or error) # 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] - - # Determine whether to use CTF or local execution - use_ctf = False - if ctf is not None: - # Check if ctf is a proper object that can handle shell commands - if (hasattr(ctf, 'get_shell') and callable(ctf.get_shell)) or isinstance(ctf, dict) or isinstance(ctf, str): - use_ctf = True - - # Run the command using the appropriate method - if use_ctf: + + # 2. Determine Execution Environment (Container > CTF > SSH > Local) + active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") + is_ssh_env = all(os.getenv(var) for var in ['SSH_USER', 'SSH_HOST']) + + # --- Docker Container Execution --- + if active_container and not ctf and not is_ssh_env: + container_id = active_container + container_workspace = _get_container_workspace_path() + context_msg = f"(docker:{container_id[:12]}:{container_workspace})" + + # Handle Async Session Creation in Container + if async_mode: + # Create a session specifically for the container environment + new_session_id = create_shell_session(command, container_id=container_id) # noqa E501 + if "Failed" in new_session_id: # Check if session creation failed + return new_session_id + if stdout: + # Wait a moment for initial output + time.sleep(0.2) + output = get_session_output(new_session_id, clear=False) + print(f"\033[32m(Started Session {new_session_id} in {context_msg})\n{output}\033[0m") # noqa E501 + return f"Started async session {new_session_id} in container {container_id[:12]}. Use this ID to interact." # noqa E501 + + # Handle Streaming Container Execution - not yet implemented for containers + if stream: + # For now, display that streaming isn't supported for containers + from cai.util import cli_print_tool_output + if call_id and tool_name: + tool_args = {"command": command, "container": container_id[:12]} + cli_print_tool_output( + tool_name, + tool_args, + "Streaming not yet supported for container execution. Running normally...", + call_id=call_id + ) + + # Handle Synchronous Execution in Container try: - # Try with CTF first - return _run_ctf(ctf, command, stdout, timeout, stream, call_id) + # Ensure container workspace exists (best effort) + # Consider moving this to workspace set/container activation + mkdir_cmd = ["docker", "exec", container_id, "mkdir", "-p", container_workspace] # noqa E501 + subprocess.run(mkdir_cmd, capture_output=True, text=True, check=False, timeout=10) # noqa E501 + + # Construct the docker exec command with workspace context + cmd_list = [ + "docker", "exec", + "-w", container_workspace, # Set working directory + container_id, + "sh", "-c", command # Execute command via shell + ] + result = subprocess.run( + cmd_list, + capture_output=True, + text=True, + check=False, # Don't raise exception on non-zero exit + timeout=timeout + ) + + output = result.stdout if result.stdout else result.stderr + output = output.strip() # Clean trailing newline + + if stdout: + print(f"\033[32m{context_msg} $ {command}\n{output}\033[0m") # noqa E501 + + # Check if command failed specifically because container isn't running + if result.returncode != 0 and "is not running" in result.stderr: + print(color(f"{context_msg} Container is not running. Attempting execution on host instead.", fg="yellow")) # noqa E501 + # Fallback to local execution, preserving workspace context + return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 + + return output # Return combined stdout/stderr + + except subprocess.TimeoutExpired: + timeout_msg = "Timeout executing command in container." + if stdout: + print(f"\033[33m{context_msg} $ {command}\nTIMEOUT\033[0m") # noqa E501 + print(color("Attempting execution on host instead.", fg="yellow")) + # Fallback to local execution on timeout + return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 except Exception as e: # pylint: disable=broad-except - # Fallback to local if CTF fails - print(color(f"CTF execution failed, falling back to local: {str(e)}", fg="yellow")) - return _run_local(command, stdout, timeout, stream, call_id, tool_name) - else: - # Use local execution - return _run_local(command, stdout, timeout, stream, call_id, tool_name) + error_msg = f"Error executing command in container: {str(e)}" + print(color(f"{context_msg} {error_msg}", fg="red")) + print(color("Attempting execution on host instead.", fg="yellow")) + # Fallback to local execution on other errors + return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 + + # --- CTF Execution --- + if ctf: + # Handling streaming for CTF - not fully implemented yet + if stream: + from cai.util import cli_print_tool_output + if call_id and tool_name: + tool_args = {"command": command, "ctf": True} + cli_print_tool_output( + tool_name, + tool_args, + "Streaming not yet supported for CTF execution. Running normally...", + call_id=call_id + ) + + # _run_ctf handles workspace internally using _get_workspace_dir() default + return _run_ctf(ctf, command, stdout, timeout) # Pass None for workspace_dir + + # --- SSH Execution --- + if is_ssh_env: + # Async for SSH would require session management via SSH client features + if async_mode: + return "Async mode not fully supported for SSH environment via this function yet." + + # Handling streaming for SSH - not fully implemented yet + if stream: + from cai.util import cli_print_tool_output + if call_id and tool_name: + tool_args = {"command": command, "ssh": True} + cli_print_tool_output( + tool_name, + tool_args, + "Streaming not yet supported for SSH execution. Running normally...", + call_id=call_id + ) + + # _run_ssh handles command execution, workspace is relative to remote home + return _run_ssh(command, stdout, timeout) # Workspace dir less relevant here + + # --- Local Execution (Default Fallback) --- + # Let _run_local handle determining the host workspace + # Handle Async Session Creation Locally + if async_mode: + # create_shell_session uses _get_workspace_dir() when container_id is None + new_session_id = create_shell_session(command) + if isinstance(new_session_id, str) and "Failed" in new_session_id: # Check failure + return new_session_id + # Retrieve the actual workspace dir the session is using + session = ACTIVE_SESSIONS.get(new_session_id) + actual_workspace = session.workspace_dir if session else "unknown" + if stdout: + time.sleep(0.2) # Allow session buffer to populate + output = get_session_output(new_session_id, clear=False) + print(f"\033[32m(Started Session {new_session_id} in local:{actual_workspace})\n{output}\033[0m") + return f"Started async session {new_session_id} locally. Use this ID to interact." + + # Handle Synchronous Execution Locally using _run_local default with streaming support + return _run_local(command, stdout, timeout, stream, call_id, tool_name, None) From 069ec7f9b1d97d8306c1227d10990ef53e14432e Mon Sep 17 00:00:00 2001 From: luijait Date: Tue, 6 May 2025 17:16:21 +0200 Subject: [PATCH 25/29] rebase common.py --- src/cai/tools/common.py | 532 ++++++++-------------------------------- 1 file changed, 98 insertions(+), 434 deletions(-) diff --git a/src/cai/tools/common.py b/src/cai/tools/common.py index e475d5f4..e1e58778 100644 --- a/src/cai/tools/common.py +++ b/src/cai/tools/common.py @@ -24,70 +24,14 @@ except ImportError: # Global dictionary to store active sessions ACTIVE_SESSIONS = {} -def _get_workspace_dir() -> str: - """Determines the target workspace directory based on env vars for host.""" - base_dir_env = os.getenv("CAI_WORKSPACE_DIR") - workspace_name = os.getenv("CAI_WORKSPACE") - - # Determine the base directory - if base_dir_env: - base_dir = os.path.abspath(base_dir_env) - else: # Default base directory is 'workspaces' - if workspace_name: - base_dir = os.path.join(os.getcwd(), "workspaces") - else: # If no workspace name is set, the workspace IS the CWD. - return os.getcwd() - - # If a workspace name is provided, append it to the base directory - if workspace_name: - if not all(c.isalnum() or c in ['_', '-'] for c in workspace_name): - print(color(f"Invalid CAI_WORKSPACE name '{workspace_name}'. " - f"Using directory '{base_dir}' instead.", fg="yellow")) - target_dir = base_dir - else: - target_dir = os.path.join(base_dir, workspace_name) - else: - target_dir = base_dir - - # Ensure the final target directory exists on the host - try: - abs_target_dir = os.path.abspath(target_dir) - os.makedirs(abs_target_dir, exist_ok=True) - return abs_target_dir - except OSError as e: - print(color(f"Error creating/accessing host workspace directory '{abs_target_dir}': {e}", - fg="red")) - print(color(f"Falling back to current directory: {os.getcwd()}", fg="yellow")) - return os.getcwd() - -def _get_container_workspace_path() -> str: - """Determines the target workspace path inside the container.""" - workspace_name = os.getenv("CAI_WORKSPACE") - if workspace_name: - if not all(c.isalnum() or c in ['_', '-'] for c in workspace_name): - print(color(f"Invalid CAI_WORKSPACE name '{workspace_name}' for container. " - f"Using '/workspace'.", fg="yellow")) - return "/" - # Standard path inside CAI containers - return f"/workspace/workspaces/{workspace_name}" - else: - return "/" class ShellSession: # pylint: disable=too-many-instance-attributes """Class to manage interactive shell sessions""" - def __init__(self, command, session_id=None, ctf=None, workspace_dir=None, container_id=None): # noqa E501 + def __init__(self, command, session_id=None, ctf=None): self.session_id = session_id or str(uuid.uuid4())[:8] - self.command = command + self.command = command self.ctf = ctf - self.container_id = container_id - # Determine workspace based on context (container, ctf or local host) - if self.container_id: - self.workspace_dir = _get_container_workspace_path() - elif self.ctf: - self.workspace_dir = workspace_dir or _get_workspace_dir() - else: - self.workspace_dir = _get_workspace_dir() self.process = None self.master = None self.slave = None @@ -96,40 +40,9 @@ class ShellSession: # pylint: disable=too-many-instance-attributes self.last_activity = time.time() def start(self): - """Start the shell session in the appropriate environment.""" - start_message_cmd = self.command - - # --- Start in Container --- - if self.container_id: - try: - self.master, self.slave = pty.openpty() - docker_cmd_list = [ - "docker", "exec", "-i", - "-w", self.workspace_dir, - self.container_id, - "sh", "-c", # Use shell to handle complex commands if needed - self.command # The actual command to run - ] - self.process = subprocess.Popen( - docker_cmd_list, - stdin=self.slave, - stdout=self.slave, - stderr=self.slave, - preexec_fn=os.setsid, - universal_newlines=True - ) - self.is_running = True - self.output_buffer.append( - f"[Session {self.session_id}] Started in container {self.container_id[:12]}: " - f"{start_message_cmd} in {self.workspace_dir}") - threading.Thread(target=self._read_output, daemon=True).start() - except Exception as e: - self.output_buffer.append(f"Error starting container session: {str(e)}") - self.is_running = False - return - - # --- Start in CTF --- + """Start the shell session""" if self.ctf: + # For CTF environments self.is_running = True self.output_buffer.append( f"[Session { @@ -139,100 +52,81 @@ class ShellSession: # pylint: disable=too-many-instance-attributes output = self.ctf.get_shell(self.command) self.output_buffer.append(output) except Exception as e: # pylint: disable=broad-except - self.output_buffer.append(f"Error executing CTF command: {str(e)}") - self.is_running = False + self.output_buffer.append(f"Error: {str(e)}") + self.is_running = False return - # --- Start Locally (Host) --- + # For local environment try: + # Create a pseudo-terminal self.master, self.slave = pty.openpty() + + # Start the process self.process = subprocess.Popen( # pylint: disable=subprocess-popen-preexec-fn, consider-using-with # noqa: E501 - self.command, - shell=True, # nosec B602 + self.command, + shell=True, # nosec B602 stdin=self.slave, stdout=self.slave, stderr=self.slave, - cwd=self.workspace_dir, - preexec_fn=os.setsid, + preexec_fn=os.setsid, # Create a new process group universal_newlines=True ) + self.is_running = True self.output_buffer.append( f"[Session { self.session_id}] Started: { self.command}") + # Start a thread to read output threading.Thread(target=self._read_output, daemon=True).start() except Exception as e: # pylint: disable=broad-except - self.output_buffer.append(f"Error starting local session: {str(e)}") + self.output_buffer.append(f"Error starting session: {str(e)}") self.is_running = False def _read_output(self): """Read output from the process""" try: - while self.is_running and self.master is not None: + while self.is_running: try: - # Check if process has exited before reading - if self.process and self.process.poll() is not None: - self.is_running = False - break - # Read the output output = os.read(self.master, 1024).decode() if output: self.output_buffer.append(output) self.last_activity = time.time() - else: + except OSError: + # No data available or terminal closed + time.sleep(0.1) + if not self.is_process_running(): self.is_running = False break - except Exception as e: - self.output_buffer.append(f"Error reading output buffer: {str(read_err)}") - self.is_running = False - break - # Add a small sleep to prevent busy-waiting if no output - if is_process_running(self): - time.sleep(0.05) - except Exception as e: - self.output_buffer.append(f"Error in read_output loop: {str(e)}") + except Exception as e: # pylint: disable=broad-except + self.output_buffer.append(f"Error reading output: {str(e)}") self.is_running = False - def is_process_running(self): """Check if the process is still running""" - # For CTF or container - if self.container_id or self.ctf: - return self.is_running - # For local host if not self.process: return False return self.process.poll() is None def send_input(self, input_data): - """Send input to the process (local or container)""" - if not self.is_running: # For CTF or container - if self.process and self.process.poll() is None: - self.is_running = True - else: # For local host - return "Session is not running" + """Send input to the process""" + if not self.is_running: + return "Session is not running" try: - # --- Send to CTF --- if self.ctf: + # For CTF environments output = self.ctf.get_shell(input_data) self.output_buffer.append(output) return "Input sent to CTF session" - # --- Send to Local or Container PTY --- - if self.master is not None: - input_data_bytes = (input_data.rstrip() + "\n").encode() - bytes_written = os.write(self.master, input_data_bytes) - if bytes_written != len(input_data_bytes): - self.output_buffer.append(f"[Session {self.session_id}] Warning: Partial input write.") - self.last_activity = time.time() - return "Input sent to session" - else: - return "Session PTY not available for input" + # For local environment + input_data = input_data.rstrip() + "\n" + os.write(self.master, input_data.encode()) + self.last_activity = time.time() + return "Input sent to session" except Exception as e: # pylint: disable=broad-except - self.output_buffer.append(f"Error sending input: {str(e)}") return f"Error sending input: {str(e)}" def get_output(self, clear=True): @@ -244,12 +138,8 @@ class ShellSession: # pylint: disable=too-many-instance-attributes def terminate(self): """Terminate the session""" - session_id_short = self.session_id[:8] if not self.is_running: - if self.process and self.process.poll() is None: - pass # Process is running, proceed with termination - else: - return f"Session {session_id_short} already terminated or finished." + return "Session already terminated" try: self.is_running = False @@ -257,64 +147,28 @@ class ShellSession: # pylint: disable=too-many-instance-attributes if self.process: # Try to terminate the process group try: - os.killpg(os.getpgid(self.process.pid), signal.SIGTERM) - except ProcessLookupError: - pass # Process already gone - except subprocess.TimeoutExpired: - print(color(f"Session {session_id_short} did not terminate gracefully, sending SIGKILL...", fg="yellow")) # noqa E501 - try: - if pgid: - os.killpg(pgid, signal.SIGKILL) # Force kill - else: - self.process.kill() - except ProcessLookupError: - pass # Already gone - except Exception as kill_err: - termination_message = f" (Error during SIGKILL: {kill_err})" - except Exception as term_err: # Catch other errors during SIGTERM - termination_message = f" (Error during SIGTERM: {term_err})" - try: - self.process.kill() - except Exception: pass # Ignore nested errors + os.killpg(os.getpgid(self.process.pid), signal.SIGTERM) + except BaseException: # pylint: disable=bare-except,broad-except # noqa: E501 + # If that fails, try to terminate just the process + self.process.terminate() + # Clean up resources + if self.master: + os.close(self.master) + if self.slave: + os.close(self.slave) - # Final check - if self.process.poll() is None: - print(color(f"Session {session_id_short} process {self.process.pid} may still be running after termination attempts.", fg="red")) # noqa E501 - termination_message += " (Warning: Process may still be running)" - - - # Clean up PTY resources if they exist - if self.master: - try: os.close(self.master) - except OSError: pass - self.master = None - if self.slave: - try: os.close(self.slave) - except OSError: pass - self.slave = None - - return termination_message or f"Session {self.session_id} terminated" + return f"Session {self.session_id} terminated" except Exception as e: # pylint: disable=broad-except - return f"Error terminating session {session_id_short}: {str(e)}" + return f"Error terminating session: {str(e)}" -def create_shell_session(command, ctf=None, container_id=None, **kwargs): - """Create a new shell session in the correct workspace/environment.""" - if container_id: - session = ShellSession(command, ctf=ctf, container_id=container_id) - else: - workspace_dir = _get_workspace_dir() - session = ShellSession(command, ctf=ctf, workspace_dir=workspace_dir) - +def create_shell_session(command, ctf=None): + """Create a new shell session""" + session = ShellSession(command, ctf=ctf) session.start() - if session.is_running or (ctf and not session.is_running): - ACTIVE_SESSIONS[session.session_id] = session - return session.session_id - else: - error_msg = session.get_output(clear=True) - print(color(f"Failed to start session: {error_msg}", fg="red")) - return f"Failed to start session: {error_msg}" + ACTIVE_SESSIONS[session.session_id] = session + return session.session_id def list_shell_sessions(): @@ -358,100 +212,47 @@ def get_session_output(session_id, clear=True): def terminate_session(session_id): """Terminate a specific session""" if session_id not in ACTIVE_SESSIONS: - return f"Session {session_id} not found or already terminated." + return f"Session {session_id} not found" session = ACTIVE_SESSIONS[session_id] result = session.terminate() - if session_id in ACTIVE_SESSIONS: - del ACTIVE_SESSIONS[session_id] + del ACTIVE_SESSIONS[session_id] return result -def _run_ctf(ctf, command, stdout=False, timeout=100, workspace_dir=None): - """Runs command in CTF env, changing to workspace_dir first.""" - target_dir = workspace_dir or _get_workspace_dir() - full_command = f"cd '{target_dir}' && {command}" - original_cmd_for_msg = command # For logging - context_msg = f"(ctf:{target_dir})" +def _run_ctf(ctf, command, stdout=False, timeout=100, stream=False, call_id=None): try: - output = ctf.get_shell(full_command, timeout=timeout) + # Ensure the command is executed in a shell that supports command + # chaining + output = ctf.get_shell(command, timeout=timeout) + # exploit_logger.log_ok() + if stdout: - print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501 - return output + print("\033[32m" + output + "\033[0m") + return output # output if output else result.stder except Exception as e: # pylint: disable=broad-except - error_msg = f"Error executing CTF command '{original_cmd_for_msg}' in '{target_dir}': {e}" # noqa E501 - print(color(error_msg, fg="red")) - return error_msg - -def _run_ssh(command, stdout=False, timeout=100, workspace_dir=None): - """Runs command via SSH. Assumes SSH agent or passwordless setup unless sshpass is used externally.""" # noqa E501 - ssh_user = os.environ.get('SSH_USER') - ssh_host = os.environ.get('SSH_HOST') - ssh_pass = os.environ.get('SSH_PASS') - remote_command = command - original_cmd_for_msg = command - context_msg = f"({ssh_user}@{ssh_host})" - - # Construct base SSH command list - if ssh_pass: - ssh_cmd_list = ["sshpass", "-p", ssh_pass, "ssh", f"{ssh_user}@{ssh_host}"] # noqa E501 - else: - ssh_cmd_list = ["ssh", f"{ssh_user}@{ssh_host}"] - ssh_cmd_list.append(remote_command) - - try: - # Use subprocess.run with list of args for better security than shell=True - result = subprocess.run( - ssh_cmd_list, - capture_output=True, - text=True, - check=False, # Don't raise exception on non-zero exit code - timeout=timeout - ) - output = result.stdout if result.stdout else result.stderr - if stdout: - print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501 - # Return combined output, potentially including errors - return output.strip() - except subprocess.TimeoutExpired as e: - error_output = e.stdout if e.stdout else str(e) - timeout_msg = f"Timeout executing SSH command: {error_output}" - if stdout: - print(f"\033[33m{context_msg} $ {original_cmd_for_msg}\nTIMEOUT\n{error_output}\033[0m") # noqa E501 - return timeout_msg - except FileNotFoundError: - # Handle case where ssh or sshpass isn't installed - error_msg = f"'sshpass' or 'ssh' command not found. Ensure they are installed and in PATH." # noqa E501 - print(color(error_msg, fg="red")) - return error_msg - except Exception as e: # pylint: disable=broad-except - error_msg = f"Error executing SSH command '{original_cmd_for_msg}' on {ssh_host}: {e}" # noqa E501 - print(color(error_msg, fg="red")) - return error_msg + print(color(f"Error executing CTF command: {e}", fg="red")) + # exploit_logger.log_error(str(e)) + return f"Error executing CTF command: {str(e)}" -def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None, workspace_dir=None): - """Runs command locally in the specified workspace_dir.""" +def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None): # If streaming is enabled and we have a call_id if stream and call_id: - return _run_local_streamed(command, call_id, timeout, tool_name, workspace_dir) + return _run_local_streamed(command, call_id, timeout, tool_name) - target_dir = workspace_dir or _get_workspace_dir() - original_cmd_for_msg = command # For logging - context_msg = f"(local:{target_dir})" try: + # nosec B602 - shell=True is required for command chaining result = subprocess.run( command, shell=True, # nosec B602 capture_output=True, text=True, - check=False, - timeout=timeout, - cwd=target_dir - ) + check=False, + timeout=timeout) output = result.stdout if result.stdout else result.stderr if stdout: - print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m") # noqa E501 + 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 @@ -460,21 +261,20 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t # Optional: Add cli_print_tool_output call here if needed for non-streaming pass - return output.strip() + return output except subprocess.TimeoutExpired as e: - error_output = e.stdout if e.stdout else str(e) + error_output = e.stdout.decode() if e.stdout else str(e) if stdout: print("\033[32m" + error_output + "\033[0m") - return error_output + return error_output except Exception as e: # pylint: disable=broad-except - error_msg = f"Error executing local command: {e}" - print(color(error_msg, fg="red")) - return error_msg + error_msg = f"Error executing local command: {e}" + print(color(error_msg, fg="red")) + return error_msg -def _run_local_streamed(command, call_id, timeout=100, tool_name=None, workspace_dir=None): - """Run a local command with streaming output to the Tool output panel.""" - target_dir = workspace_dir or _get_workspace_dir() +def _run_local_streamed(command, call_id, timeout=100, tool_name=None): + """Run a local command with streaming output to the Tool output panel""" try: # Try to import Rich for nice display try: @@ -498,8 +298,7 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None, workspace stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, - bufsize=1, - cwd=target_dir # Set CWD for local process + bufsize=1 ) # If tool_name is not provided, derive it from the command @@ -529,7 +328,7 @@ def _run_local_streamed(command, call_id, timeout=100, tool_name=None, workspace header.append(")", style="yellow") tool_time = 0 start_time = time.time() - total_time = time.time() - START_TIME + total_time = time.time() - START_TIME timing_info = [] if total_time: timing_info.append(f"Total: {format_time(total_time)}") @@ -706,8 +505,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg async_mode=False, session_id=None, timeout=100, stream=False, call_id=None, tool_name=None): """ - Run command in the appropriate environment (Docker, CTF, SSH, Local) - and workspace. + Run command either in CTF container or on the local attacker machine Args: command: The command to execute @@ -722,168 +520,34 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg If None, the tool name will be derived from the command. Returns: - str: Command output, status message, or session ID. + str: Command output, status message, or session ID """ # If session_id is provided, send command to that session if session_id: if session_id not in ACTIVE_SESSIONS: return f"Session {session_id} not found" - session = ACTIVE_SESSIONS[session_id] - result = session.send_input(command) # Send the raw command string + + result = send_to_session(session_id, command) if stdout: output = get_session_output(session_id, clear=False) - env_type = "Local" - if session.container_id: - env_type = f"Container({session.container_id[:12]})" - elif session.ctf: - env_type = "CTF" - print(f"\033[32m(Session {session_id} in {env_type}:{session.workspace_dir}) >> {command}\n{output}\033[0m") # noqa E501 - return result # Return the result of sending input ("Input sent..." or error) + print("\033[32m" + output + "\033[0m") + return result + + # If async_mode, create a new session + if async_mode: + session_id = create_shell_session(command, ctf) + if stdout: + # Wait a moment for initial output + 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." # 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] - - # 2. Determine Execution Environment (Container > CTF > SSH > Local) - active_container = os.getenv("CAI_ACTIVE_CONTAINER", "") - is_ssh_env = all(os.getenv(var) for var in ['SSH_USER', 'SSH_HOST']) - - # --- Docker Container Execution --- - if active_container and not ctf and not is_ssh_env: - container_id = active_container - container_workspace = _get_container_workspace_path() - context_msg = f"(docker:{container_id[:12]}:{container_workspace})" - - # Handle Async Session Creation in Container - if async_mode: - # Create a session specifically for the container environment - new_session_id = create_shell_session(command, container_id=container_id) # noqa E501 - if "Failed" in new_session_id: # Check if session creation failed - return new_session_id - if stdout: - # Wait a moment for initial output - time.sleep(0.2) - output = get_session_output(new_session_id, clear=False) - print(f"\033[32m(Started Session {new_session_id} in {context_msg})\n{output}\033[0m") # noqa E501 - return f"Started async session {new_session_id} in container {container_id[:12]}. Use this ID to interact." # noqa E501 - - # Handle Streaming Container Execution - not yet implemented for containers - if stream: - # For now, display that streaming isn't supported for containers - from cai.util import cli_print_tool_output - if call_id and tool_name: - tool_args = {"command": command, "container": container_id[:12]} - cli_print_tool_output( - tool_name, - tool_args, - "Streaming not yet supported for container execution. Running normally...", - call_id=call_id - ) - - # Handle Synchronous Execution in Container - try: - # Ensure container workspace exists (best effort) - # Consider moving this to workspace set/container activation - mkdir_cmd = ["docker", "exec", container_id, "mkdir", "-p", container_workspace] # noqa E501 - subprocess.run(mkdir_cmd, capture_output=True, text=True, check=False, timeout=10) # noqa E501 - - # Construct the docker exec command with workspace context - cmd_list = [ - "docker", "exec", - "-w", container_workspace, # Set working directory - container_id, - "sh", "-c", command # Execute command via shell - ] - result = subprocess.run( - cmd_list, - capture_output=True, - text=True, - check=False, # Don't raise exception on non-zero exit - timeout=timeout - ) - - output = result.stdout if result.stdout else result.stderr - output = output.strip() # Clean trailing newline - - if stdout: - print(f"\033[32m{context_msg} $ {command}\n{output}\033[0m") # noqa E501 - - # Check if command failed specifically because container isn't running - if result.returncode != 0 and "is not running" in result.stderr: - print(color(f"{context_msg} Container is not running. Attempting execution on host instead.", fg="yellow")) # noqa E501 - # Fallback to local execution, preserving workspace context - return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 - - return output # Return combined stdout/stderr - - except subprocess.TimeoutExpired: - timeout_msg = "Timeout executing command in container." - if stdout: - print(f"\033[33m{context_msg} $ {command}\nTIMEOUT\033[0m") # noqa E501 - print(color("Attempting execution on host instead.", fg="yellow")) - # Fallback to local execution on timeout - return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 - except Exception as e: # pylint: disable=broad-except - error_msg = f"Error executing command in container: {str(e)}" - print(color(f"{context_msg} {error_msg}", fg="red")) - print(color("Attempting execution on host instead.", fg="yellow")) - # Fallback to local execution on other errors - return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501 - - # --- CTF Execution --- + + # Otherwise, run command normally if ctf: - # Handling streaming for CTF - not fully implemented yet - if stream: - from cai.util import cli_print_tool_output - if call_id and tool_name: - tool_args = {"command": command, "ctf": True} - cli_print_tool_output( - tool_name, - tool_args, - "Streaming not yet supported for CTF execution. Running normally...", - call_id=call_id - ) - - # _run_ctf handles workspace internally using _get_workspace_dir() default - return _run_ctf(ctf, command, stdout, timeout) # Pass None for workspace_dir - - # --- SSH Execution --- - if is_ssh_env: - # Async for SSH would require session management via SSH client features - if async_mode: - return "Async mode not fully supported for SSH environment via this function yet." - - # Handling streaming for SSH - not fully implemented yet - if stream: - from cai.util import cli_print_tool_output - if call_id and tool_name: - tool_args = {"command": command, "ssh": True} - cli_print_tool_output( - tool_name, - tool_args, - "Streaming not yet supported for SSH execution. Running normally...", - call_id=call_id - ) - - # _run_ssh handles command execution, workspace is relative to remote home - return _run_ssh(command, stdout, timeout) # Workspace dir less relevant here - - # --- Local Execution (Default Fallback) --- - # Let _run_local handle determining the host workspace - # Handle Async Session Creation Locally - if async_mode: - # create_shell_session uses _get_workspace_dir() when container_id is None - new_session_id = create_shell_session(command) - if isinstance(new_session_id, str) and "Failed" in new_session_id: # Check failure - return new_session_id - # Retrieve the actual workspace dir the session is using - session = ACTIVE_SESSIONS.get(new_session_id) - actual_workspace = session.workspace_dir if session else "unknown" - if stdout: - time.sleep(0.2) # Allow session buffer to populate - output = get_session_output(new_session_id, clear=False) - print(f"\033[32m(Started Session {new_session_id} in local:{actual_workspace})\n{output}\033[0m") - return f"Started async session {new_session_id} locally. Use this ID to interact." - - # Handle Synchronous Execution Locally using _run_local default with streaming support - return _run_local(command, stdout, timeout, stream, call_id, tool_name, None) + return _run_ctf(ctf, command, stdout, timeout, stream, call_id) + return _run_local(command, stdout, timeout, stream, call_id, tool_name) From 23e1e65052ee3eab34ef89af6bc23b987df957d2 Mon Sep 17 00:00:00 2001 From: luijait Date: Tue, 6 May 2025 17:18:29 +0200 Subject: [PATCH 26/29] Delete unnecesary lines on run.py --- src/cai/sdk/agents/run.py | 33 --------------------------------- 1 file changed, 33 deletions(-) diff --git a/src/cai/sdk/agents/run.py b/src/cai/sdk/agents/run.py index 30d82d2f..e2379560 100644 --- a/src/cai/sdk/agents/run.py +++ b/src/cai/sdk/agents/run.py @@ -779,22 +779,6 @@ class Runner: run_config: RunConfig, tool_use_tracker: AgentToolUseTracker, ) -> SingleStepResult: - # Log the raw model response output, focusing on tool calls - logger.debug(f"[_get_single_step_result_from_response] Raw new_response.id: {new_response.referenceable_id}") - if new_response.output: - for i, item in enumerate(new_response.output): - item_type = type(item).__name__ - logger.debug(f"[_get_single_step_result_from_response] Raw output item [{i}] type: {item_type}") - if hasattr(item, "name") and hasattr(item, "arguments"): # For ResponseFunctionToolCall - logger.debug( - f"[_get_single_step_result_from_response] Raw ResponseFunctionToolCall item [{i}]: " - f"Name='{getattr(item, 'name', 'N/A')}', " - f"Args='{getattr(item, 'arguments', 'N/A')}', " - f"CallID='{getattr(item, 'call_id', 'N/A')}'" - ) - elif hasattr(item, "text"): # For ResponseOutputText - logger.debug(f"[_get_single_step_result_from_response] Raw ResponseOutputText item [{i}]: Text='{getattr(item, 'text', '')[:100]}...'") - processed_response = RunImpl.process_model_response( agent=agent, @@ -803,26 +787,9 @@ class Runner: output_schema=output_schema, handoffs=handoffs, ) - - # Log the processed response, focusing on tools_used - logger.debug(f"[_get_single_step_result_from_response] Processed response type: {type(processed_response).__name__}") - if hasattr(processed_response, 'is_final_output'): - logger.debug(f"[_get_single_step_result_from_response] Processed response: is_final_output={processed_response.is_final_output}") - else: - logger.debug(f"[_get_single_step_result_from_response] Processed response does not have is_final_output attribute") - - if hasattr(processed_response, 'final_output_from_llm') and processed_response.final_output_from_llm is not None: - logger.debug(f"[_get_single_step_result_from_response] Processed response: final_output_from_llm='{str(processed_response.final_output_from_llm)[:100]}...'") # Log tools used with robust type checking if hasattr(processed_response, 'tools_used') and processed_response.tools_used: - # Log summarizing number of tools used first - logger.debug(f"[_get_single_step_result_from_response] Found {len(processed_response.tools_used)} tools used") - - # Add spacing between blocks in terminal output - if os.environ.get('CAI_STREAM', 'false').lower() == 'true': - print("") # Visual separator only when streaming is enabled - for i, tool_call in enumerate(processed_response.tools_used): try: # Safely extract tool name with multiple fallbacks From d67f944b5e6af17cb764de906b51afb0474ac2c0 Mon Sep 17 00:00:00 2001 From: luijait Date: Tue, 6 May 2025 17:19:32 +0200 Subject: [PATCH 27/29] Delete unnecesary lines on run.py --- src/cai/sdk/agents/run.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/cai/sdk/agents/run.py b/src/cai/sdk/agents/run.py index e2379560..103661e0 100644 --- a/src/cai/sdk/agents/run.py +++ b/src/cai/sdk/agents/run.py @@ -4,7 +4,6 @@ import asyncio import copy from dataclasses import dataclass, field from typing import Any, cast -import os from openai.types.responses import ResponseCompletedEvent From 3fb20717217db71610713c5c77fd6328114388f3 Mon Sep 17 00:00:00 2001 From: luijait Date: Tue, 6 May 2025 17:22:02 +0200 Subject: [PATCH 28/29] Delete trash on openaichatcompletions --- src/cai/sdk/agents/models/openai_chatcompletions.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index 21c5ef81..bf35c06c 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -1089,10 +1089,6 @@ class OpenAIChatCompletionsModel(Model): direct_stats = final_stats.copy() direct_stats["interaction_cost"] = float(interaction_cost) direct_stats["total_cost"] = float(total_cost) - - # Add a small delay to avoid overlapping outputs - await asyncio.sleep(0.3) - # Use the direct copy with guaranteed float costs finish_agent_streaming(streaming_context, direct_stats) @@ -1103,12 +1099,6 @@ class OpenAIChatCompletionsModel(Model): # Find the assistant message to print for item in final_response.output: if isinstance(item, ResponseOutputMessage) and item.role == 'assistant': - # Add a small delay to avoid overlapping outputs - await asyncio.sleep(0.3) - - # Print clear visual separator before message - print("\n") - cli_print_agent_messages( agent_name=getattr(self, 'agent_name', 'Agent'), message=item, @@ -1612,7 +1602,6 @@ class OpenAIChatCompletionsModel(Model): } }] - # Handle special Ollama generic_linux_command format if isinstance(delta, dict) and 'content' in delta: content = delta['content'] # Try to detect if the content is a JSON string with function call format From e6c9b5bfe8ab22b09c011b42eebd506545eeb3c1 Mon Sep 17 00:00:00 2001 From: luijait Date: Tue, 6 May 2025 17:24:33 +0200 Subject: [PATCH 29/29] Delete trash on openaichatcompletions --- src/cai/sdk/agents/models/openai_chatcompletions.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/cai/sdk/agents/models/openai_chatcompletions.py b/src/cai/sdk/agents/models/openai_chatcompletions.py index bf35c06c..18134f83 100644 --- a/src/cai/sdk/agents/models/openai_chatcompletions.py +++ b/src/cai/sdk/agents/models/openai_chatcompletions.py @@ -550,10 +550,6 @@ class OpenAIChatCompletionsModel(Model): model_str = str(self.model).lower() is_ollama = self.is_ollama or "ollama" in model_str or ":" in model_str or "qwen" in model_str - - # Add a small delay to make sure any previous tool outputs are fully rendered - # This helps prevent overlapping of panels in the terminal - await asyncio.sleep(0.2) # Add visual separation before agent output if streaming_context and should_show_rich_stream: