diff --git a/textgen.py b/textgen.py
index 6a41a58..ab55553 100644
--- a/textgen.py
+++ b/textgen.py
@@ -35,8 +35,13 @@ def train_ai():
 match sys.argv[1]:
     case "generate":
         # ? send message to parent JS process
-        print(generate_message(sys.argv[2]))
-        sys.stdout.flush()
+        try:
+            prompt = ' '.join(map(str, sys.argv[2:]))
+            print(generate_message(prompt))
+        except IndexError:
+            print(generate_message())
+        finally:
+            sys.stdout.flush()
     case "train":
         train_ai()
     case _: