Fix headless detection

README.md

@@ -52,6 +52,7 @@ resp = api1.send_message('Hello, world!')
 print(resp['message'])
 
 api1.reset_conversation()  # reset the conversation
+api1.close()  # close the session
 ```
 
 ## Frequently Asked Questions

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "pyChatGPT"
-version = "0.3.1"
+version = "0.3.2"
 authors = [
     { name="terry3041", email="[email protected]" },
 ]

src/pyChatGPT/__main__.py

@@ -14,17 +14,20 @@ def clear_screen():
 
     clear_screen()
     print(
-        'Conversation started. Type "reset" to reset the conversation. Type "reauth" to reauthenticate.'
+        'Conversation started. Type "reset" to reset the conversation. Type "quit" to quit.'
     )
     while True:
         prompt = input('\nYou: ')
         if prompt.lower() == 'reset':
             chat.reset_conversation()
             clear_screen()
             print(
-                'Conversation started. Type "reset" to reset the conversation. Type "reauth" to reauthenticate.'
+                'Conversation started. Type "reset" to reset the conversation. Type "quit" to quit.'
             )
             continue
+        if prompt.lower() == 'quit':
+            chat.close()
+            break
         print('\nChatGPT: ', end='')
         response = chat.send_message(prompt)
         print(response['message'])

src/pyChatGPT/pyChatGPT.py

@@ -7,8 +7,9 @@
 
 from pyvirtualdisplay import Display
 import markdownify
-import re
+import platform
 import os
+import re
 
 
 class ChatGPT:
@@ -33,9 +34,18 @@ def __init__(
             raise ValueError('Invalid proxy format')
 
         self.session_token = session_token
-        self.is_headless = os.name == 'posix' and 'DISPLAY' not in os.environ
+        self.is_headless = platform.system() == 'Linux' and 'DISPLAY' not in os.environ
         self.__init_browser()
 
+    def close(self) -> None:
+        '''
+        Close the browser and stop the virtual display (if any)
+        '''
+        if hasattr(self, 'driver'):
+            self.driver.quit()
+        if hasattr(self, 'display'):
+            self.display.stop()
+
     def __init_browser(self) -> None:
         '''
         Get the Cloudflare cookies & user-agent