mirror of https://github.com/aliasrobotics/cai.git
Streamv2beta
This commit is contained in:
luijait
parent
c16d901498
commit
82eb56409a
|
|
@ -0,0 +1 @@
|
|||
print('Hola mundo')
|
||||
|
|
@ -355,7 +355,7 @@ def run_cai_cli(starting_agent, context_variables=None, stream=False, max_turns=
|
|||
title="[bold]Session Summary[/bold]",
|
||||
title_align="left"
|
||||
)
|
||||
console.print(time_panel)
|
||||
console.print(time_panel, end="")
|
||||
|
||||
print_session_summary(console, metrics, logging_path)
|
||||
|
||||
|
|
|
|||
|
|
@ -238,7 +238,7 @@ class AgentCommand(Command):
|
|||
os.environ["CAI_AGENT_TYPE"] = selected_agent_key
|
||||
|
||||
console.print(
|
||||
f"[green]Switched to agent: {agent_name}[/green]")
|
||||
f"[green]Switched to agent: {agent_name}[/green]", end="")
|
||||
visualize_agent_graph(agent)
|
||||
return True
|
||||
|
||||
|
|
|
|||
|
|
@ -421,7 +421,7 @@ class ModelCommand(Command):
|
|||
change_message,
|
||||
border_style="green",
|
||||
title="Model Changed"
|
||||
)
|
||||
), end=""
|
||||
)
|
||||
return True
|
||||
|
||||
|
|
|
|||
|
|
@ -177,7 +177,7 @@ def display_banner(console: Console):
|
|||
[white] Bug bounty-ready AI[/white]
|
||||
"""
|
||||
|
||||
console.print(banner)
|
||||
console.print(banner, end="")
|
||||
|
||||
# # Create a table showcasing CAI framework capabilities
|
||||
# #
|
||||
|
|
@ -361,4 +361,4 @@ def display_quick_guide(console: Console):
|
|||
border_style="blue",
|
||||
padding=(1, 2),
|
||||
title_align="center"
|
||||
))
|
||||
), end="")
|
||||
|
|
|
|||
|
|
@ -96,7 +96,7 @@ def get_user_input(
|
|||
|
||||
# Get user input with all features
|
||||
return prompt(
|
||||
[('class:prompt', '\nCAI> ')],
|
||||
[('class:prompt', 'CAI> ')],
|
||||
completer=command_completer,
|
||||
style=create_prompt_style(),
|
||||
history=FileHistory(str(history_file)),
|
||||
|
|
|
|||
|
|
@ -434,7 +434,7 @@ def _run_ssh(command, stdout=False, timeout=100, workspace_dir=None):
|
|||
return error_msg
|
||||
|
||||
|
||||
def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None, workspace_dir=None):
|
||||
def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, tool_name=None, workspace_dir=None, custom_args=None):
|
||||
"""Runs command locally in the specified workspace_dir."""
|
||||
# Make sure we're in active time mode for tool execution
|
||||
stop_idle_timer()
|
||||
|
|
@ -453,17 +453,17 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t
|
|||
|
||||
# Parse command into parts for display
|
||||
parts = command.strip().split(' ', 1)
|
||||
cmd = parts[0] if parts else ""
|
||||
cmd_var = parts[0] if parts else ""
|
||||
args_param_val = parts[1] if len(parts) > 1 else "" # Renamed to avoid conflict with tool_args dict key
|
||||
|
||||
# For generic Linux commands, standardize the tool_name format
|
||||
if not tool_name:
|
||||
tool_name = f"{cmd}_command" if cmd else "command"
|
||||
tool_name = f"{cmd_var}_command" if cmd_var else "command"
|
||||
|
||||
# Create args dictionary with non-empty values only
|
||||
tool_args = {}
|
||||
if cmd:
|
||||
tool_args["command"] = cmd
|
||||
if cmd_var:
|
||||
tool_args["command"] = cmd_var
|
||||
if args_param_val and args_param_val.strip():
|
||||
tool_args["args"] = args_param_val
|
||||
|
||||
|
|
@ -471,9 +471,16 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t
|
|||
tool_args["workspace"] = os.path.basename(target_dir)
|
||||
tool_args["full_command"] = command
|
||||
|
||||
# If custom args were provided, merge them with the default args
|
||||
if custom_args is not None:
|
||||
if isinstance(custom_args, dict):
|
||||
# Merge the dictionaries, with custom args taking precedence
|
||||
for key, value in custom_args.items():
|
||||
tool_args[key] = value
|
||||
|
||||
# For generic commands, ensure we have a unique call_id
|
||||
if not call_id:
|
||||
call_id = f"cmd_{cmd}_{str(uuid.uuid4())[:8]}"
|
||||
call_id = f"cmd_{cmd_var}_{str(uuid.uuid4())[:8]}"
|
||||
|
||||
# Initialize/use the call_id for this streaming session
|
||||
call_id = start_tool_streaming(tool_name, tool_args, call_id)
|
||||
|
|
@ -561,11 +568,11 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t
|
|||
# If this is the same command that was recently streamed, don't display it again
|
||||
# We'll create a similar tool_args structure to what streaming would use
|
||||
parts = command.strip().split(' ', 1)
|
||||
cmd = parts[0] if parts else ""
|
||||
cmd_var = parts[0] if parts else ""
|
||||
args_param_val = parts[1] if len(parts) > 1 else ""
|
||||
|
||||
# Generate a standard tool name for consistency with streaming
|
||||
standard_tool_name = tool_name or (f"{cmd}_command" if cmd else "command")
|
||||
standard_tool_name = tool_name or (f"{cmd_var}_command" if cmd_var else "command")
|
||||
|
||||
# Calculate a consistent command key - must match the format used in cli_print_tool_output
|
||||
command_key = f"{standard_tool_name}:{args_param_val}"
|
||||
|
|
@ -575,7 +582,37 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t
|
|||
return output.strip()
|
||||
|
||||
if stdout:
|
||||
print(f"\033[32m{context_msg} $ {original_cmd_for_msg}\n{output}\033[0m")
|
||||
# Create a tool_args dictionary for non-streaming display
|
||||
# that matches the format used in streaming
|
||||
tool_display_args = {
|
||||
"command": cmd_var,
|
||||
"args": args_param_val,
|
||||
"full_command": command,
|
||||
"workspace": os.path.basename(target_dir)
|
||||
}
|
||||
|
||||
# If custom args were provided, merge them with the default args
|
||||
if custom_args is not None and isinstance(custom_args, dict):
|
||||
for key, value in custom_args.items():
|
||||
tool_display_args[key] = value
|
||||
|
||||
# Calculate execution info
|
||||
tool_execution_time = time.time() - process_start_time
|
||||
exec_info = {
|
||||
"status": "completed" if result.returncode == 0 else "error",
|
||||
"return_code": result.returncode,
|
||||
"environment": "Local",
|
||||
"host": os.path.basename(target_dir),
|
||||
"tool_time": tool_execution_time
|
||||
}
|
||||
|
||||
# Display the command output with rich formatting
|
||||
cli_print_tool_output(
|
||||
tool_name=standard_tool_name,
|
||||
args=tool_display_args,
|
||||
output=output,
|
||||
execution_info=exec_info
|
||||
)
|
||||
|
||||
return output.strip()
|
||||
except subprocess.TimeoutExpired as e:
|
||||
|
|
@ -587,13 +624,13 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t
|
|||
from cai.util import finish_tool_streaming
|
||||
# Parse the command the same way we did for streaming
|
||||
parts = command.strip().split(' ', 1)
|
||||
cmd = parts[0] if parts else ""
|
||||
args = parts[1] if len(parts) > 1 else ""
|
||||
cmd_var = parts[0] if parts else ""
|
||||
args_var = parts[1] if len(parts) > 1 else ""
|
||||
|
||||
# Ensure tool_args has complete information
|
||||
tool_args = {
|
||||
"command": cmd,
|
||||
"args": args if args.strip() else "",
|
||||
"command": cmd_var,
|
||||
"args": args_var if args_var.strip() else "",
|
||||
"full_command": command,
|
||||
"environment": "Local",
|
||||
"workspace": os.path.basename(target_dir)
|
||||
|
|
@ -604,7 +641,7 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t
|
|||
"environment": "Local",
|
||||
"host": os.path.basename(target_dir)
|
||||
}
|
||||
finish_tool_streaming(tool_name or f"{cmd}_command", tool_args, error_msg, call_id, execution_info)
|
||||
finish_tool_streaming(tool_name or f"{cmd_var}_command", tool_args, error_msg, call_id, execution_info)
|
||||
|
||||
if stdout:
|
||||
print("\033[32m" + error_msg + "\033[0m")
|
||||
|
|
@ -618,13 +655,13 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t
|
|||
from cai.util import finish_tool_streaming
|
||||
# Parse the command the same way we did for streaming
|
||||
parts = command.strip().split(' ', 1)
|
||||
cmd = parts[0] if parts else ""
|
||||
args = parts[1] if len(parts) > 1 else ""
|
||||
cmd_var = parts[0] if parts else ""
|
||||
args_var = parts[1] if len(parts) > 1 else ""
|
||||
|
||||
# Ensure tool_args has complete information
|
||||
tool_args = {
|
||||
"command": cmd,
|
||||
"args": args if args.strip() else "",
|
||||
"command": cmd_var,
|
||||
"args": args_var if args_var.strip() else "",
|
||||
"full_command": command,
|
||||
"environment": "Local",
|
||||
"workspace": os.path.basename(target_dir)
|
||||
|
|
@ -635,7 +672,7 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t
|
|||
"environment": "Local",
|
||||
"host": os.path.basename(target_dir)
|
||||
}
|
||||
finish_tool_streaming(tool_name or f"{cmd}_command", tool_args, error_msg, call_id, execution_info)
|
||||
finish_tool_streaming(tool_name or f"{cmd_var}_command", tool_args, error_msg, call_id, execution_info)
|
||||
|
||||
print(color(error_msg, fg="red"))
|
||||
return error_msg
|
||||
|
|
@ -647,7 +684,7 @@ def _run_local(command, stdout=False, timeout=100, stream=False, call_id=None, t
|
|||
|
||||
def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arguments # noqa: E501
|
||||
async_mode=False, session_id=None,
|
||||
timeout=100, stream=False, call_id=None, tool_name=None):
|
||||
timeout=100, stream=False, call_id=None, tool_name=None, args=None):
|
||||
"""
|
||||
Run command in the appropriate environment (Docker, CTF, SSH, Local)
|
||||
and workspace.
|
||||
|
|
@ -663,6 +700,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
|
|||
call_id: Unique ID for the command execution (for streaming)
|
||||
tool_name: Name of the tool being executed (for display in streaming output).
|
||||
If None, the tool name will be derived from the command.
|
||||
args: Additional arguments for the tool (for display and context).
|
||||
|
||||
Returns:
|
||||
str: Command output, status message, or session ID.
|
||||
|
|
@ -673,17 +711,17 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
|
|||
|
||||
# Parse command into standard parts to ensure consistent naming
|
||||
parts = command.strip().split(' ', 1)
|
||||
cmd = parts[0] if parts else ""
|
||||
args = parts[1] if len(parts) > 1 else ""
|
||||
cmd_name = parts[0] if parts else ""
|
||||
cmd_args = parts[1] if len(parts) > 1 else ""
|
||||
|
||||
# Generate a call_id if we're streaming and one wasn't provided
|
||||
# Use a more specific format that includes the command name for easier tracking
|
||||
if not call_id and stream:
|
||||
call_id = f"cmd_{cmd}_{str(uuid.uuid4())[:8]}"
|
||||
call_id = f"cmd_{cmd_name}_{str(uuid.uuid4())[:8]}"
|
||||
|
||||
# If no tool_name is provided, derive it from the command in a consistent way
|
||||
if not tool_name:
|
||||
tool_name = f"{cmd}_command" if cmd else "command"
|
||||
tool_name = f"{cmd_name}_command" if cmd_name else "command"
|
||||
|
||||
try:
|
||||
# If session_id is provided, send command to that session
|
||||
|
|
@ -750,8 +788,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
|
|||
|
||||
# Create args dictionary with standardized format
|
||||
tool_args = {
|
||||
"command": cmd,
|
||||
"args": args if args.strip() else "",
|
||||
"command": cmd_name,
|
||||
"args": cmd_args if cmd_args.strip() else "",
|
||||
"full_command": command,
|
||||
"container": container_id[:12],
|
||||
"environment": "Container",
|
||||
|
|
@ -887,7 +925,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
|
|||
start_idle_timer()
|
||||
# Fallback to local execution on timeout
|
||||
print(color("Container execution timed out. Attempting execution on host instead.", fg="yellow"))
|
||||
return _run_local(command, stdout, timeout, False, None, tool_name, _get_workspace_dir())
|
||||
return _run_local(command, stdout, timeout, False, None, tool_name, _get_workspace_dir(), args)
|
||||
|
||||
except Exception as e:
|
||||
# Handle other errors
|
||||
|
|
@ -908,7 +946,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
|
|||
start_idle_timer()
|
||||
# Fallback to local execution on error
|
||||
print(color("Container execution failed. Attempting execution on host instead.", fg="yellow"))
|
||||
return _run_local(command, stdout, timeout, False, None, tool_name, _get_workspace_dir())
|
||||
return _run_local(command, stdout, timeout, False, None, tool_name, _get_workspace_dir(), args)
|
||||
|
||||
# Handle Synchronous Execution in Container
|
||||
try:
|
||||
|
|
@ -945,7 +983,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
|
|||
stop_active_timer()
|
||||
start_idle_timer()
|
||||
# Fallback to local execution, preserving workspace context
|
||||
return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501
|
||||
return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir(), args) # noqa E501
|
||||
|
||||
# Switch back to idle mode after command completes
|
||||
stop_active_timer()
|
||||
|
|
@ -961,7 +999,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
|
|||
stop_active_timer()
|
||||
start_idle_timer()
|
||||
# Fallback to local execution on timeout
|
||||
return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501
|
||||
return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir(), args) # 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"))
|
||||
|
|
@ -970,7 +1008,7 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
|
|||
stop_active_timer()
|
||||
start_idle_timer()
|
||||
# Fallback to local execution on other errors
|
||||
return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir()) # noqa E501
|
||||
return _run_local(command, stdout, timeout, stream, call_id, tool_name, _get_workspace_dir(), args) # noqa E501
|
||||
|
||||
# --- CTF Execution ---
|
||||
if ctf:
|
||||
|
|
@ -981,8 +1019,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
|
|||
|
||||
# Create args dictionary with standardized format
|
||||
tool_args = {
|
||||
"command": cmd,
|
||||
"args": args if args.strip() else "",
|
||||
"command": cmd_name,
|
||||
"args": cmd_args if cmd_args.strip() else "",
|
||||
"full_command": command,
|
||||
"environment": "CTF",
|
||||
"workspace": os.path.basename(_get_workspace_dir())
|
||||
|
|
@ -1066,8 +1104,8 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
|
|||
|
||||
# Create args dictionary with standardized format
|
||||
tool_args = {
|
||||
"command": cmd,
|
||||
"args": args if args.strip() else "",
|
||||
"command": cmd_name,
|
||||
"args": cmd_args if cmd_args.strip() else "",
|
||||
"full_command": command,
|
||||
"ssh_host": ssh_connection,
|
||||
"environment": "SSH"
|
||||
|
|
@ -1213,7 +1251,9 @@ def run_command(command, ctf=None, stdout=False, # pylint: disable=too-many-arg
|
|||
timeout,
|
||||
stream=stream,
|
||||
call_id=call_id,
|
||||
tool_name=tool_name
|
||||
tool_name=tool_name,
|
||||
workspace_dir=_get_workspace_dir(),
|
||||
custom_args=args
|
||||
)
|
||||
|
||||
# Switch back to idle mode after local command completes
|
||||
|
|
|
|||
|
|
@ -37,62 +37,101 @@ def execute_code(code: str = "", language: str = "python",
|
|||
"python": "py",
|
||||
"php": "php",
|
||||
"bash": "sh",
|
||||
"shell": "sh", # Add shell as alias for bash
|
||||
"ruby": "rb",
|
||||
"perl": "pl",
|
||||
"golang": "go",
|
||||
"go": "go", # Add go as alias for golang
|
||||
"javascript": "js",
|
||||
"js": "js", # Add js as alias for javascript
|
||||
"typescript": "ts",
|
||||
"ts": "ts", # Add ts as alias for typescript
|
||||
"rust": "rs",
|
||||
"csharp": "cs",
|
||||
"cs": "cs", # Add cs as alias for csharp
|
||||
"java": "java",
|
||||
"kotlin": "kt"
|
||||
"kotlin": "kt",
|
||||
"c": "c", # Add C language
|
||||
"cpp": "cpp", # Add C++ language
|
||||
"c++": "cpp" # Add C++ language alias
|
||||
}
|
||||
ext = extensions.get(language.lower(), "txt")
|
||||
# Normalize language to lowercase
|
||||
language = language.lower()
|
||||
ext = extensions.get(language, "txt")
|
||||
full_filename = f"{filename}.{ext}"
|
||||
|
||||
# Create code file with content
|
||||
create_cmd = f"cat << 'EOF' > {full_filename}\n{code}\nEOF"
|
||||
result = run_command(create_cmd, ctf=ctf)
|
||||
result = run_command(create_cmd, ctf=ctf, stream=True, tool_name="execute_code")
|
||||
if "error" in result.lower():
|
||||
return f"Failed to create code file: {result}"
|
||||
if language.lower() == "python":
|
||||
|
||||
# Prepare execution command based on language
|
||||
if language in ["python", "py"]:
|
||||
exec_cmd = f"python3 {full_filename}"
|
||||
elif language.lower() == "php":
|
||||
elif language in ["php"]:
|
||||
exec_cmd = f"php {full_filename}"
|
||||
elif language.lower() in ["bash", "sh"]:
|
||||
elif language in ["bash", "sh", "shell"]:
|
||||
exec_cmd = f"bash {full_filename}"
|
||||
elif language.lower() == "ruby":
|
||||
elif language in ["ruby", "rb"]:
|
||||
exec_cmd = f"ruby {full_filename}"
|
||||
elif language.lower() == "perl":
|
||||
elif language in ["perl", "pl"]:
|
||||
exec_cmd = f"perl {full_filename}"
|
||||
elif language.lower() == "golang" or language.lower() == "go":
|
||||
elif language in ["golang", "go"]:
|
||||
temp_dir = f"/tmp/go_exec_{filename}"
|
||||
run_command(f"mkdir -p {temp_dir}", ctf=ctf)
|
||||
run_command(f"cp {full_filename} {temp_dir}/main.go", ctf=ctf)
|
||||
run_command(f"cd {temp_dir} && go mod init temp", ctf=ctf)
|
||||
exec_cmd = f"cd {temp_dir} && go run main.go"
|
||||
elif language.lower() == "javascript":
|
||||
elif language in ["javascript", "js"]:
|
||||
exec_cmd = f"node {full_filename}"
|
||||
elif language.lower() == "typescript":
|
||||
elif language in ["typescript", "ts"]:
|
||||
exec_cmd = f"ts-node {full_filename}"
|
||||
elif language.lower() == "rust":
|
||||
elif language in ["rust", "rs"]:
|
||||
# For Rust, we need to compile first
|
||||
run_command(f"rustc {full_filename} -o {filename}", ctf=ctf)
|
||||
exec_cmd = f"./{filename}"
|
||||
elif language.lower() == "csharp":
|
||||
elif language in ["csharp", "cs"]:
|
||||
# For C#, compile with dotnet
|
||||
run_command(f"dotnet build {full_filename}", ctf=ctf)
|
||||
exec_cmd = f"dotnet run {full_filename}"
|
||||
elif language.lower() == "java":
|
||||
elif language in ["java"]:
|
||||
# For Java, compile first
|
||||
run_command(f"javac {full_filename}", ctf=ctf)
|
||||
exec_cmd = f"java {filename}"
|
||||
elif language.lower() == "kotlin":
|
||||
elif language in ["kotlin", "kt"]:
|
||||
# For Kotlin, compile first
|
||||
run_command(f"kotlinc {full_filename} -include-runtime -d {filename}.jar", ctf=ctf)
|
||||
exec_cmd = f"java -jar {filename}.jar"
|
||||
elif language in ["c"]:
|
||||
# For C, compile with gcc
|
||||
run_command(f"gcc {full_filename} -o {filename}", ctf=ctf)
|
||||
exec_cmd = f"./{filename}"
|
||||
elif language in ["cpp", "c++"]:
|
||||
# For C++, compile with g++
|
||||
run_command(f"g++ {full_filename} -o {filename}", ctf=ctf)
|
||||
exec_cmd = f"./{filename}"
|
||||
else:
|
||||
return f"Unsupported language: {language}"
|
||||
|
||||
output = run_command(exec_cmd, ctf=ctf, timeout=timeout)
|
||||
# Execute the code with syntax-highlighted output
|
||||
# Create a custom tool args dictionary to send language and code info to the tool output function
|
||||
tool_args = {
|
||||
"command": "execute",
|
||||
"language": language,
|
||||
"filename": filename,
|
||||
"code": code, # Include the code for syntax highlighting
|
||||
"timeout": timeout
|
||||
}
|
||||
|
||||
# Run the command with streaming to get syntax highlighting
|
||||
output = run_command(
|
||||
exec_cmd,
|
||||
ctf=ctf,
|
||||
timeout=timeout,
|
||||
stream=True,
|
||||
tool_name="execute_code",
|
||||
args=tool_args
|
||||
)
|
||||
|
||||
return output
|
||||
|
|
|
|||
139
src/cai/util.py
139
src/cai/util.py
|
|
@ -22,6 +22,10 @@ from dataclasses import dataclass, field
|
|||
from typing import Dict, Optional
|
||||
import time
|
||||
import threading
|
||||
from rich.syntax import Syntax # Import Syntax for highlighting
|
||||
from rich.panel import Panel
|
||||
from rich.console import Group
|
||||
from rich.box import ROUNDED
|
||||
|
||||
# Global timing variables for tracking active and idle time
|
||||
_active_timer_start = None
|
||||
|
|
@ -1183,7 +1187,7 @@ def create_agent_streaming_context(agent_name, counter, model):
|
|||
Text.assemble(header, content, footer),
|
||||
border_style="blue",
|
||||
box=ROUNDED,
|
||||
padding=(1, 2),
|
||||
padding=(0, 1),
|
||||
title="[bold]Agent Streaming Response[/bold]",
|
||||
title_align="left",
|
||||
width=panel_width,
|
||||
|
|
@ -1244,7 +1248,7 @@ def update_agent_streaming_content(context, text_delta):
|
|||
Text.assemble(context["header"], context["content"], context["footer"]),
|
||||
border_style="blue",
|
||||
box=ROUNDED,
|
||||
padding=(1, 2),
|
||||
padding=(0, 1),
|
||||
title="[bold]Agent Streaming Response[/bold]",
|
||||
title_align="left",
|
||||
width=context.get("panel_width", 100),
|
||||
|
|
@ -1349,13 +1353,12 @@ def finish_agent_streaming(context, final_stats=None):
|
|||
Text.assemble(
|
||||
context["header"],
|
||||
context["content"],
|
||||
Text("\n\n"),
|
||||
tokens_text if tokens_text else Text(""),
|
||||
tokens_text if tokens_text else Text(""),
|
||||
context["footer"]
|
||||
),
|
||||
border_style="blue",
|
||||
box=ROUNDED,
|
||||
padding=(1, 2),
|
||||
padding=(0, 1),
|
||||
title="[bold]Agent Streaming Response[/bold]",
|
||||
title_align="left",
|
||||
width=context.get("panel_width", 100),
|
||||
|
|
@ -1522,7 +1525,7 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut
|
|||
content,
|
||||
title=title,
|
||||
border_style=border_style,
|
||||
padding=(1, 2),
|
||||
padding=(0, 1),
|
||||
box=ROUNDED,
|
||||
title_align="left"
|
||||
)
|
||||
|
|
@ -1586,18 +1589,34 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut
|
|||
# Create a console for output
|
||||
console = Console(theme=theme)
|
||||
|
||||
# Get the panel content
|
||||
# Get the panel content - with syntax highlighting
|
||||
header, content = _create_tool_panel_content(tool_name, args, output, execution_info, token_info)
|
||||
|
||||
# Determine border style
|
||||
# Determine border style based on status
|
||||
border_style = "blue" # Default for non-streaming
|
||||
|
||||
if execution_info:
|
||||
status = execution_info.get('status', 'completed')
|
||||
if status == "completed":
|
||||
border_style = "green"
|
||||
title = "[bold green]Tool Output [Completed][/bold green]"
|
||||
elif status == "error":
|
||||
border_style = "red"
|
||||
title = "[bold red]Tool Output [Error][/bold red]"
|
||||
elif status == "timeout":
|
||||
border_style = "red"
|
||||
title = "[bold red]Tool Output [Timeout][/bold red]"
|
||||
else:
|
||||
title = "[bold blue]Tool Output[/bold blue]"
|
||||
else:
|
||||
title = "[bold blue]Tool Output[/bold blue]"
|
||||
|
||||
# Create the panel
|
||||
panel = Panel(
|
||||
content,
|
||||
title="[bold blue]Tool Output[/bold blue]",
|
||||
title=title,
|
||||
border_style=border_style,
|
||||
padding=(1, 2),
|
||||
padding=(0, 1),
|
||||
box=ROUNDED,
|
||||
title_align="left"
|
||||
)
|
||||
|
|
@ -1613,6 +1632,10 @@ def cli_print_tool_output(tool_name="", args="", output="", call_id=None, execut
|
|||
def _create_tool_panel_content(tool_name, args, output, execution_info=None, token_info=None):
|
||||
"""Create the header and content for a tool output panel."""
|
||||
from rich.text import Text
|
||||
from rich.syntax import Syntax # Import Syntax for highlighting
|
||||
from rich.panel import Panel
|
||||
from rich.console import Group
|
||||
from rich.box import ROUNDED
|
||||
|
||||
# Format arguments for display
|
||||
args_str = _format_tool_args(args)
|
||||
|
|
@ -1655,7 +1678,101 @@ def _create_tool_panel_content(tool_name, args, output, execution_info=None, tok
|
|||
# Create token information if available
|
||||
token_content = _create_token_info_display(token_info)
|
||||
|
||||
# Assemble the full content
|
||||
# Determine if we need specialized content formatting
|
||||
group_content = [header]
|
||||
|
||||
# Special handling for execute_code tool
|
||||
if tool_name == "execute_code":
|
||||
try:
|
||||
# Parse args to get language and code
|
||||
language = "python" # Default
|
||||
code = ""
|
||||
filename = "exploit"
|
||||
|
||||
if isinstance(args, dict):
|
||||
language = args.get("language", "python")
|
||||
code = args.get("code", "")
|
||||
filename = args.get("filename", "exploit")
|
||||
elif isinstance(args, str):
|
||||
# Try to extract language and code from args string
|
||||
import re
|
||||
lang_match = re.search(r"language\s*=\s*['\"]?([a-zA-Z0-9+]+)['\"]?", args)
|
||||
if lang_match:
|
||||
language = lang_match.group(1)
|
||||
|
||||
# Create syntax highlighted code panel if we have code
|
||||
if code:
|
||||
syntax = Syntax(code, language, theme="monokai", line_numbers=True,
|
||||
background_color="#272822", indent_guides=True)
|
||||
code_panel = Panel(
|
||||
syntax,
|
||||
title=f"Code ({language})",
|
||||
border_style="cyan",
|
||||
title_align="left",
|
||||
box=ROUNDED,
|
||||
padding=(0, 1)
|
||||
)
|
||||
|
||||
# Create output panel with proper highlighting
|
||||
output_syntax = Syntax(output, "text", theme="monokai",
|
||||
background_color="#272822")
|
||||
output_panel = Panel(
|
||||
output_syntax,
|
||||
title="Output",
|
||||
border_style="green",
|
||||
title_align="left",
|
||||
box=ROUNDED,
|
||||
padding=(0, 1)
|
||||
)
|
||||
|
||||
# Add code and output panels to content
|
||||
group_content.append("\n")
|
||||
group_content.append(code_panel)
|
||||
group_content.append("\n")
|
||||
group_content.append(output_panel)
|
||||
|
||||
# Add token info if available
|
||||
if token_content:
|
||||
group_content.append("\n")
|
||||
group_content.append(token_content)
|
||||
|
||||
return header, Group(*group_content)
|
||||
except Exception:
|
||||
# Fallback if syntax highlighting fails
|
||||
pass
|
||||
|
||||
# Special handling for generic_linux_command
|
||||
elif tool_name == "generic_linux_command" or "command" in tool_name:
|
||||
try:
|
||||
# Highlight the output as bash
|
||||
output_syntax = Syntax(output, "bash", theme="monokai",
|
||||
background_color="#272822")
|
||||
|
||||
# Create a panel for the formatted output
|
||||
output_panel = Panel(
|
||||
output_syntax,
|
||||
title="Command Output",
|
||||
border_style="green",
|
||||
title_align="left",
|
||||
box=ROUNDED,
|
||||
padding=(0, 1)
|
||||
)
|
||||
|
||||
# Assemble content with highlighted output
|
||||
group_content.append("\n")
|
||||
group_content.append(output_panel)
|
||||
|
||||
# Add token info if available
|
||||
if token_content:
|
||||
group_content.append("\n")
|
||||
group_content.append(token_content)
|
||||
|
||||
return header, Group(*group_content)
|
||||
except Exception:
|
||||
# Fallback if syntax highlighting fails
|
||||
pass
|
||||
|
||||
# Default content assembly for other tools
|
||||
content = Text()
|
||||
content.append(header)
|
||||
content.append("\n\n")
|
||||
|
|
|
|||
Loading…
Reference in New Issue