From 4c1fe834a889c5aa29d339e5b1e198ec91a8f020 Mon Sep 17 00:00:00 2001
From: LeonOstrez <leon.ostrez@gmail.com>
Date: Fri, 29 Sep 2023 13:31:11 +0100
Subject: [PATCH] fix getting args.root, ask feedback only on success and
 keyboard interrupt (not on exceptions), fix terminal width and erasing old
 logs

---
 pilot/logger/logger.py        | 11 ++++++++---
 pilot/main.py                 |  2 +-
 pilot/utils/exit.py           |  6 ++++--
 pilot/utils/files.py          |  2 +-
 pilot/utils/llm_connection.py |  5 ++++-
 5 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/pilot/logger/logger.py b/pilot/logger/logger.py
index 3e19c7895..603cb7ace 100644
--- a/pilot/logger/logger.py
+++ b/pilot/logger/logger.py
@@ -31,14 +31,19 @@ def setup_logger():
 
 
 def filter_sensitive_fields(record):
-    if len(record.args):
+    if isinstance(record.args, dict):  # check if args is a dictionary
         args = record.args.copy()
-
         for field in sensitive_fields:
             if field in args:
                 args[field] = '*****'
-
         record.args = args
+
+    elif isinstance(record.args, tuple):  # check if args is a tuple
+        args_list = list(record.args)
+        # Convert the tuple to a list and replace sensitive fields
+        args_list = ['*****' if arg in sensitive_fields else arg for arg in args_list]
+        record.args = tuple(args_list)
+
     return record.levelno <= logging.INFO
 
 
diff --git a/pilot/main.py b/pilot/main.py
index 027df7aaf..8f46e15d2 100644
--- a/pilot/main.py
+++ b/pilot/main.py
@@ -91,6 +91,6 @@ def local_print(*args, **kwargs):
         print(red('---------- GPT PILOT EXITING WITH ERROR ----------'))
         traceback.print_exc()
         print(red('--------------------------------------------------'))
-        exit_gpt_pilot()
+        exit_gpt_pilot(False)
     finally:
         sys.exit(0)
diff --git a/pilot/utils/exit.py b/pilot/utils/exit.py
index 708419c39..4b5d286dd 100644
--- a/pilot/utils/exit.py
+++ b/pilot/utils/exit.py
@@ -42,10 +42,12 @@ def get_path_id():
     return hashlib.sha256(installation_directory.encode()).hexdigest()
 
 
-def exit_gpt_pilot():
+def exit_gpt_pilot(ask_feedback=True):
     path_id = get_path_id()
     send_telemetry(path_id)
 
-    feedback = get_user_feedback()
+    feedback = None
+    if ask_feedback:
+        feedback = get_user_feedback()
     if feedback:  # only send if user provided feedback
         send_feedback(feedback, path_id)
diff --git a/pilot/utils/files.py b/pilot/utils/files.py
index 6536e9c4a..e065a6fed 100644
--- a/pilot/utils/files.py
+++ b/pilot/utils/files.py
@@ -20,7 +20,7 @@ def setup_workspace(args):
 
         return args['workspace']
 
-    root = args['root'] or get_parent_folder('pilot')
+    root = args.get('root') or get_parent_folder('pilot')
     create_directory(root, 'workspace')
     project_path = create_directory(os.path.join(root, 'workspace'), args['name'])
     create_directory(project_path, 'tests')
diff --git a/pilot/utils/llm_connection.py b/pilot/utils/llm_connection.py
index c5db88ade..d683029f0 100644
--- a/pilot/utils/llm_connection.py
+++ b/pilot/utils/llm_connection.py
@@ -200,7 +200,10 @@ def stream_gpt_completion(data, req_type):
     """
 
     # TODO add type dynamically - this isn't working when connected to the external process
-    terminal_width = 50  # os.get_terminal_size().columns
+    try:
+        terminal_width = os.get_terminal_size().columns
+    except OSError:
+        terminal_width = 50
     lines_printed = 2
     gpt_response = ''
     buffer = ''  # A buffer to accumulate incoming data