Add timeout and enhance console UI

WebChatGPT/console.py

@@ -56,7 +56,6 @@ def start_spinning(
         except Exception as e:
             cls.querying = False
             logging.debug(getExc(e))
-        finally:
             t1.join()
 
     @classmethod
@@ -94,9 +93,9 @@ class InteractiveChatGPT(cmd.Cmd):
         f"┌─[{getpass.getuser().capitalize()}@WebChatGPT]({__version__})\r\n└──╼ ❯❯❯"
     )
 
-    def __init__(self, auth, cookie_path, model, index, *args, **kwargs):
+    def __init__(self, auth, cookie_path, model, index, timeout, *args, **kwargs):
         super().__init__(*args, **kwargs)
-        self.bot = ChatGPT(auth, cookie_path, model, index)
+        self.bot = ChatGPT(auth, cookie_path, model, index, timeout=timeout)
 
     def do_help(self, text):
         """Echoes useful help info
@@ -146,7 +145,10 @@ def default(self, line):
                 busy_bar.start_spinning()
                 generated_response = self.bot.chat(line)
                 busy_bar.stop_spinning()
-                rich.print(Markdown(generated_response))
+                if self.prettify:
+                    rich.print(Markdown(generated_response))
+                else:
+                    click.secho(generated_response)
 
             except (KeyboardInterrupt, EOFError):
                 busy_bar.stop_spinning()
@@ -190,6 +192,13 @@ def chat():
 @click.option(
     "-I", "--index", help="Conversation index to resume from", type=click.INT, default=0
 )
+@click.option(
+    "-T",
+    "--timeout",
+    help="Http request timeout",
+    type=click.INT,
+    default=30,
+)
 @click.option(
     "-P",
     "--prompt",
@@ -203,11 +212,15 @@ def chat():
     default=1,
     envvar="busy_bar_index",
 )
-def interactive(auth, cookie_path, model, index, prompt, busy_bar_index):
+@click.option("--prettify/--raw", default=True, help="Prettify the markdowned response")
+def interactive(
+    auth, cookie_path, model, index, timeout, prompt, busy_bar_index, prettify
+):
     """Chat with ChatGPT interactively"""
     assert isinstance(busy_bar_index, int), "Index must be an integer only"
     busy_bar.spin_index = busy_bar_index
-    bot = InteractiveChatGPT(auth, cookie_path, model, index)
+    bot = InteractiveChatGPT(auth, cookie_path, model, index, timeout)
+    bot.prettify = prettify
     if prompt:
         bot.default(prompt)
     bot.cmdloop()
@@ -239,18 +252,29 @@ def interactive(auth, cookie_path, model, index, prompt, busy_bar_index):
 @click.option(
     "-I", "--index", help="Conversation index to resume from", type=click.INT, default=0
 )
+@click.option(
+    "-T",
+    "--timeout",
+    help="Http request timeout",
+    type=click.INT,
+    default=30,
+)
 @click.option(
     "-P",
     "--prompt",
     help="Start conversation with this messsage",
     prompt="Enter message",
 )
-def generate(auth, cookie_path, model, index, prompt):
+@click.option("--prettify/--raw", default=True, help="Prettify the markdowned response")
+def generate(auth, cookie_path, model, index, timeout, prompt, prettify):
     """Generate a quick response with ChatGPT"""
 
-    content = ChatGPT(auth, cookie_path, model, index).chat(prompt)
+    content = ChatGPT(auth, cookie_path, model, index, timeout=timeout).chat(prompt)
 
-    rich.print(Markdown(content))
+    if prettify:
+        rich.print(Markdown(content))
+    else:
+        click.secho(content)
 
 
 @error_handler(exit_on_error=True)

WebChatGPT/main.py

@@ -15,18 +15,22 @@ def __init__(
         conversation_index: int = 0,
         locale: str = "en-US",
         user_agent: str = "Mozilla/5.0 (X11; Linux x86_64; rv:120.0) Gecko/20100101 Firefox/120.0",
+        timeout: tuple = 30,
     ):
         """Initializes ChatGPT
 
         Args:
             authorization (str): OpenAI's authorization value
             cookie_path (str): Path to `.json` file containing `chat.openai.com` cookies
-            model (str, optional): ChatGPT text generation model name. Defaults to "text-davinci-002-render-sha".
-            conversation_index (int, optional): Conversation index to pick up conversation from. Defaults to 0.
-            locale (str, optional): Your locale. Defaults to en-US
-            user_agent (str, optional): Http request header User-Agent. Defaults to Mozilla/5.0 (X11; Linux x86_64; rv:120.0) Gecko/20100101 Firefox/120.0
+            model (str, optional): ChatGPT text generation model name. Defaults to `text-davinci-002-render-sha`.
+            conversation_index (int, optional): Conversation index to pick up conversation from. Defaults to `0`.
+            locale (str, optional): Your locale. Defaults to `en-US`
+            user_agent (str, optional): Http request header User-Agent. Defaults to `Mozilla/5.0 (X11; Linux x86_64; rv:120.0) Gecko/20100101 Firefox/120.0`
+            timeout (int, optional): Http request timeout.
+
         """
         self.session = requests.Session()
+        self.timeout = timeout
         self.session.headers.update(utils.get_request_headers(authorization))
         self.session.cookies.update(utils.get_cookies(cookie_path))
         self.conversation_endpoint = "https://chat.openai.com/backend-api/conversation"
@@ -107,7 +111,9 @@ def ask(self, prompt: str, stream: bool = False) -> dict:
         ```
         """
         response = self.session.post(
-            url=self.conversation_endpoint, json=self.__generate_payload(prompt)
+            url=self.conversation_endpoint,
+            json=self.__generate_payload(prompt),
+            timeout=self.timeout,
         )
         if (
             response.ok
@@ -306,7 +312,8 @@ def user_details(self, in_details: bool = True) -> dict:
         resp = self.session.get(
             self.account_details_endpoint
             if in_details
-            else self.account_detail_endpoint
+            else self.account_detail_endpoint,
+            timeout=self.timeout,
         )
         return utils.is_json(resp, "account data")
 
@@ -345,6 +352,7 @@ def prompt_library(self, limit: int = 4, offset: int = 0) -> list:
         resp = self.session.get(
             self.prompt_library_endpoint,
             params={"limit": limit, "offset": offset},
+            timeout=self.timeout,
         )
         return utils.is_json(resp, "prompts")
 
@@ -407,6 +415,7 @@ def previous_conversations(
         resp = self.session.get(
             self.previous_conversations_endpoint,
             params={"limit": limit, "offset": offset, "order": order},
+            timeout=self.timeout,
         )
         resp = utils.is_json(resp, "conversation history")
         if all:
@@ -433,5 +442,6 @@ def generate_title(self, conversation_id: str, message_id: str) -> dict:
         resp = self.session.post(
             self.title_generation_endpoint % {"conversation_id": conversation_id},
             json={"message_id": message_id},
+            timeout=self.timeout,
         )
         return utils.is_json(resp, "title")

docs/CHANGELOG.md

@@ -6,4 +6,5 @@ Initial Release
 
 **Whats New?**
 
-- General code improvement
+- General code improvement
+- Request timeout - 30