246 lines
9.1 KiB
Python
246 lines
9.1 KiB
Python
"""
|
|
Console UI components for the Edison application.
|
|
"""
|
|
import logging
|
|
import pyperclip
|
|
from termcolor import colored
|
|
from edison.core import api_client, command_executor
|
|
from edison.utils import os_utils, validation, markdown_utils
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
def handle_streaming_output(token):
|
|
"""
|
|
Handle a token from the streaming API response.
|
|
|
|
Args:
|
|
token (str): A token from the streaming response.
|
|
"""
|
|
# Print the token without a newline and flush immediately
|
|
print(token, end='', flush=True)
|
|
|
|
def print_command(command, config=None):
|
|
"""
|
|
Print a command to the console.
|
|
|
|
Args:
|
|
command (str): The command to print.
|
|
config (dict, optional): Configuration dictionary.
|
|
"""
|
|
if config and config.get("ui", {}).get("rich_formatting", True):
|
|
command_style = config.get("ui", {}).get("command_style", "panel")
|
|
theme = config.get("ui", {}).get("theme", "monokai")
|
|
shell = config.get("shell", "bash")
|
|
|
|
if command_style == "panel":
|
|
# Use rich panel style
|
|
from edison.utils.markdown_utils import print_command_rich
|
|
print_command_rich(command, shell=shell, theme=theme)
|
|
else:
|
|
# Use simple highlighting
|
|
print("Command: " + colored(command, 'blue'))
|
|
else:
|
|
# Use traditional styling
|
|
print("Command: " + colored(command, 'blue'))
|
|
|
|
def prompt_user_input(config, response):
|
|
"""
|
|
Prompt the user for input on what to do with the generated command.
|
|
|
|
Args:
|
|
config (dict): The configuration dictionary.
|
|
response (str): The generated command.
|
|
|
|
Returns:
|
|
str: The user's input.
|
|
"""
|
|
print_command(response, config)
|
|
|
|
if config.get("safety", True):
|
|
prompt_text = "Execute command? [Y]es [n]o [m]odify (prompt/command) [c]opy to clipboard ==> "
|
|
|
|
if os_utils.missing_posix_display():
|
|
prompt_text = "Execute command? [Y]es [n]o [m]odify (prompt/command) ==> "
|
|
|
|
print(prompt_text, end='')
|
|
user_input = input()
|
|
else:
|
|
user_input = "Y"
|
|
|
|
return user_input
|
|
|
|
def handle_command_execution_streaming(client, config, query, explain=False):
|
|
"""
|
|
Handle the execution of a command with streaming output.
|
|
|
|
Args:
|
|
client (OpenAI): The OpenAI client.
|
|
config (dict): The configuration dictionary.
|
|
query (str): The user's query.
|
|
explain (bool): Whether to explain the command.
|
|
"""
|
|
# Check if we should show the "Generating command: " prefix
|
|
if config.get("show_generating_prefix", True):
|
|
print(colored("Generating command: ", 'yellow'), end='', flush=True)
|
|
|
|
try:
|
|
# Call the streaming API
|
|
command = api_client.generate_command_streaming(client, config, query, handle_streaming_output)
|
|
|
|
# Print a newline after streaming is complete
|
|
print()
|
|
|
|
# Continue with the regular command execution flow
|
|
handle_command_execution(client, config, command, explain)
|
|
except KeyboardInterrupt:
|
|
print("\nCommand generation cancelled.")
|
|
return
|
|
except Exception as e:
|
|
logger.error(f"Error generating command: {e}")
|
|
print(colored(f"\nError generating command: {e}", 'red'))
|
|
|
|
def handle_command_execution(client, config, command, explain=False):
|
|
"""
|
|
Handle the execution of a command based on user input.
|
|
|
|
Args:
|
|
client (OpenAI): The OpenAI client.
|
|
config (dict): The configuration dictionary.
|
|
command (str): The command to execute.
|
|
explain (bool): Whether to explain the command.
|
|
"""
|
|
# Check for issues in the response
|
|
if validation.check_for_issue(command):
|
|
print(colored("There was an issue: " + command, 'red'))
|
|
return
|
|
|
|
# Check for markdown in the response
|
|
if validation.check_for_markdown(command):
|
|
print(colored(
|
|
"The proposed command contains markdown, response not executed directly: \n", 'red'
|
|
) + command)
|
|
return
|
|
|
|
# Get explanation if requested
|
|
if explain:
|
|
try:
|
|
from edison.ui.interactive import get_command_explanation
|
|
|
|
# Get UI configuration
|
|
ui_config = config.get("ui", {})
|
|
structured = ui_config.get("structured_explanations", True)
|
|
use_rich = ui_config.get("rich_formatting", True)
|
|
|
|
# Generate explanation
|
|
explanation = get_command_explanation(client, command, structured=structured)
|
|
|
|
if use_rich:
|
|
print() # Add a blank line before the explanation
|
|
# Format the explanation with rich formatting
|
|
formatted_explanation = markdown_utils.format_command_explanation(explanation, use_rich=True)
|
|
print(formatted_explanation)
|
|
else:
|
|
print(colored("\nExplanation:", "green", attrs=["bold"]))
|
|
# Format the explanation with terminal-friendly markdown
|
|
formatted_explanation = markdown_utils.format_command_explanation(explanation)
|
|
print(formatted_explanation)
|
|
|
|
print()
|
|
except Exception as e:
|
|
logger.error(f"Error generating explanation: {e}")
|
|
print(colored(f"Error generating explanation: {e}", 'red'))
|
|
|
|
# Get user input
|
|
user_input = prompt_user_input(config, command)
|
|
print()
|
|
|
|
# Handle user input
|
|
handle_user_input(client, config, user_input, command)
|
|
|
|
def handle_user_input(client, config, user_input, command):
|
|
"""
|
|
Handle user input for command execution.
|
|
|
|
Args:
|
|
client (OpenAI): The OpenAI client.
|
|
config (dict): The configuration dictionary.
|
|
user_input (str): The user's input.
|
|
command (str): The command to execute.
|
|
"""
|
|
if user_input.upper() == "Y" or user_input == "":
|
|
try:
|
|
command_executor.execute_command(config.get("shell", os_utils.get_default_shell()), command)
|
|
except Exception as e:
|
|
logger.error(f"Error executing command: {e}")
|
|
print(colored(f"Error executing command: {e}", 'red'))
|
|
|
|
elif user_input.upper() == "M":
|
|
print("Modify [p]rompt or [c]ommand? [c] ==> ", end='')
|
|
mod_choice = input().lower()
|
|
|
|
if mod_choice in ["p", "prompt"]:
|
|
# Modify the prompt (natural language description)
|
|
print("Modify prompt: ", end='')
|
|
modded_query = input()
|
|
|
|
if not modded_query.strip():
|
|
print(colored("Empty prompt. Command execution cancelled.", "yellow"))
|
|
return
|
|
|
|
# Use streaming if enabled in config
|
|
if config.get("streaming", True):
|
|
try:
|
|
handle_command_execution_streaming(client, config, modded_query)
|
|
except Exception as e:
|
|
logger.error(f"Error generating modified command: {e}")
|
|
print(colored(f"Error generating modified command: {e}", 'red'))
|
|
else:
|
|
try:
|
|
modded_response = api_client.generate_command(client, config, modded_query)
|
|
handle_command_execution(client, config, modded_response)
|
|
except Exception as e:
|
|
logger.error(f"Error generating modified command: {e}")
|
|
print(colored(f"Error generating modified command: {e}", 'red'))
|
|
else:
|
|
# Modify the command directly
|
|
print("Modify command: ", end='')
|
|
# Pre-fill with current command
|
|
import readline
|
|
readline.set_startup_hook(lambda: readline.insert_text(command))
|
|
try:
|
|
modded_command = input()
|
|
finally:
|
|
readline.set_startup_hook()
|
|
|
|
if not modded_command.strip():
|
|
print(colored("Empty command. Command execution cancelled.", "yellow"))
|
|
return
|
|
|
|
# Print the modified command with styling
|
|
print_command(modded_command, config)
|
|
|
|
# Ask for confirmation
|
|
print("Execute modified command? [Y/n] ==> ", end='')
|
|
confirm = input().lower()
|
|
|
|
if confirm in ["n", "no"]:
|
|
print(colored("Command execution cancelled.", "yellow"))
|
|
return
|
|
|
|
# Execute the modified command
|
|
try:
|
|
command_executor.execute_command(config.get("shell", os_utils.get_default_shell()), modded_command)
|
|
except Exception as e:
|
|
logger.error(f"Error executing modified command: {e}")
|
|
print(colored(f"Error executing modified command: {e}", 'red'))
|
|
|
|
elif user_input.upper() == "C":
|
|
if os_utils.missing_posix_display():
|
|
return
|
|
try:
|
|
pyperclip.copy(command)
|
|
print("Copied command to clipboard.")
|
|
except Exception as e:
|
|
logger.error(f"Error copying to clipboard: {e}")
|
|
print(colored(f"Error copying to clipboard: {e}", 'red'))
|