From 8c61170f5a541a3a0c81fb41738dfb7435cce694 Mon Sep 17 00:00:00 2001
From: jiacai2050 <dev@liujiacai.net>
Date: Sun, 16 Jun 2024 21:23:26 +0800
Subject: [PATCH] refactor: remove shell actions

---
 shellgpt/app.py          | 46 ++++++++++++++++++----------------------
 shellgpt/utils/common.py |  5 +++--
 2 files changed, 24 insertions(+), 27 deletions(-)

diff --git a/shellgpt/app.py b/shellgpt/app.py
index f0f913d..e8ba8b6 100644
--- a/shellgpt/app.py
+++ b/shellgpt/app.py
@@ -57,11 +57,18 @@ def tui(self, history, initial_prompt):
     def repl_action(self, prompt):
         if 'exit' == prompt:
             sys.exit(0)
-
-        if prompt.upper() == 'C':
+        elif prompt == 'c':
             copy_text(self.last_answer)
             return True
 
+        if self.is_shell:
+            if prompt == 'e':
+                self.explain_cmd(self.last_answer)
+                return True
+            elif prompt == 'r':
+                print(execute_cmd(self.last_answer, ask=True))
+                return True
+
         # Following parse set command
         if prompt.startswith('set') is False:
             return False
@@ -93,8 +100,9 @@ def repl(self, initial_prompt):
  \ \/\/ / -_) / _/ _ \ '  \/ -_) |  _/ _ \ \__ \ ' \/ -_) | | (_ |  _/ | |
   \_/\_/\___|_\__\___/_|_|_\___|  \__\___/ |___/_||_\___|_|_|\___|_|   |_|
 
-Type "exit" to exit; "c" to copy last answer
-""")
+Type "exit" to exit; "c" to copy last answer.
+When system is shell , type "e" to explain, "r" to run last command.
+""", end='')
         self.infer(initial_prompt)
         while True:
             prompt = input(f'{self.llm.system_content}@{self.llm.model}> ')
@@ -123,30 +131,18 @@ def infer(self, prompt):
                 print()
 
             self.last_answer = buf
-            if self.is_shell:
-                self.shell_action(buf)
         except Exception as e:
             print(f'Error when infer: ${e}')
 
-    def shell_action(self, cmd):
-        if IS_TTY:
-            action = input('(R)un, (C)opy, (E)xplain> ')
-            action = action.upper()
-            buf = ''
-            if action == 'E':
-                for r in self.llm.chat(
-                    f'Explain this command: {cmd}', add_system_message=False
-                ):
-                    buf += r
-                    print(r, end='', flush=True)
-                print()
-                self.last_answer = buf
-            elif action == 'R':
-                print(execute_cmd(cmd))
-            elif action == 'C':
-                copy_text(cmd)
-            else:
-                self.infer(action)
+    def explain_cmd(self, cmd):
+        buf = ''
+        for r in self.llm.chat(
+            f'Explain this command: {cmd}', add_system_message=False
+        ):
+            buf += r
+            print(r, end='', flush=True)
+        print()
+        self.last_answer = buf
 
 
 def load_system_content_when_necessary(sc):
diff --git a/shellgpt/utils/common.py b/shellgpt/utils/common.py
index eaf99b6..da1b753 100644
--- a/shellgpt/utils/common.py
+++ b/shellgpt/utils/common.py
@@ -57,8 +57,9 @@ def copy_text(text):
     pyperclip.copy(text)
 
 
-def execute_cmd(cmd):
-    return subprocess.getoutput(cmd)
+def execute_cmd(cmd, ask=False):
+    if 'y' == input(f'Type y to run: `{cmd}`> '):
+        return subprocess.getoutput(cmd)
 
 
 def base64_image(image_path: str) -> str: