diff --git a/yolo.py b/yolo.py
index 961a036..305700c 100755
--- a/yolo.py
+++ b/yolo.py
@@ -274,7 +274,6 @@ def missing_posix_display():
 
     return display == b'\n'
 
-# TODO: Change the output according the configuration
 def prompt_user_input(config, response):
     """
     Print the command proposal in blue and prompt the user for next action based on the safety
@@ -297,10 +296,15 @@ def prompt_user_input(config, response):
     print("Command: " + colored(response, 'blue'))
 
     if config["safety"]:
-        prompt_text = "Execute command? [Y]es [n]o [m]odify [c]opy to clipboard ==> "
+        modify_text = ""
+
+        if config["modify"]:
+            modify_text = " [m]modify"
+        
+        prompt_text =  "Execute command? [Y]es [n]o" + modify_text + ": "
 
         if os.name == "posix" and missing_posix_display():
-            prompt_text =  "Execute command? [Y]es [n]o [m]odify ==> "
+            prompt_text =  "Execute command? [Y]es [n]o" + modify_text + ": "
 
         print(prompt_text, end = '')
 
@@ -310,7 +314,7 @@ def prompt_user_input(config, response):
 
     return user_input
 
-def evaluate_input(config, user_input, command):
+def evaluate_input(client, config, user_input, command):
     """
     Evaluate the user input to either execute, modify, or copy the command.
 
@@ -338,16 +342,15 @@ def evaluate_input(config, user_input, command):
             # Unix: /bin/bash /bin/zsh: uses -c both Ubuntu and macOS should work, others might not
             subprocess.run([config["shell"], "-c", command], shell=False, check=True)
 
-    # TODO: change call_open_ai function
     if user_input.upper() == "M":
         print("Modify prompt: ", end = '')
         modded_query = input()
-        modded_response = call_open_ai(config, modded_query)
+        modded_response = chat_completion(client, modded_query, config)
         check_for_issue(modded_response)
         check_for_markdown(modded_response)
         modded_user_input = prompt_user_input(config, modded_response)
         print()
-        evaluate_input(config, modded_user_input, modded_response)
+        evaluate_input(client, config, modded_user_input, modded_response)
 
     if user_input.upper() == "C":
         if os.name == "posix" and missing_posix_display():
@@ -397,7 +400,7 @@ def main():
 
     user_input = prompt_user_input(config, result)
     print()
-    evaluate_input(config, user_input, result)
+    evaluate_input(client, config, user_input, result)
 
 if __name__ == "__main__":
     main()