From ab49de8b434036d7e130da54b1c9cbea6a54aa16 Mon Sep 17 00:00:00 2001 From: Simatwa Date: Sun, 17 Dec 2023 15:02:56 +0300 Subject: [PATCH] Add timeout and enhance console UI --- WebChatGPT/console.py | 42 +++++++++++++++++++++++++++++++++--------- WebChatGPT/main.py | 22 ++++++++++++++++------ docs/CHANGELOG.md | 3 ++- 3 files changed, 51 insertions(+), 16 deletions(-) diff --git a/WebChatGPT/console.py b/WebChatGPT/console.py index db9cfa9..d1cb727 100644 --- a/WebChatGPT/console.py +++ b/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) diff --git a/WebChatGPT/main.py b/WebChatGPT/main.py index 9eb804e..5663df4 100755 --- a/WebChatGPT/main.py +++ b/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") diff --git a/docs/CHANGELOG.md b/docs/CHANGELOG.md index b671bb8..9e02ea8 100644 --- a/docs/CHANGELOG.md +++ b/docs/CHANGELOG.md @@ -6,4 +6,5 @@ Initial Release **Whats New?** -- General code improvement \ No newline at end of file +- General code improvement +- Request timeout - 30 \ No newline at end of file