From 425d6ac0bf6496c73331bdaffc93e69ec61cd8c2 Mon Sep 17 00:00:00 2001 From: Nicholas Albion Date: Sat, 9 Sep 2023 10:43:59 +1000 Subject: [PATCH] spinner for LLM timeout --- pilot/utils/llm_connection.py | 6 ++++++ pilot/utils/spinner.py | 3 ++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/pilot/utils/llm_connection.py b/pilot/utils/llm_connection.py index 97b9e23..17a7b64 100644 --- a/pilot/utils/llm_connection.py +++ b/pilot/utils/llm_connection.py @@ -139,8 +139,11 @@ def count_lines_based_on_width(content, width): def retry_on_exception(func): def wrapper(*args, **kwargs): + spinner = None + while True: try: + spinner_stop(spinner) return func(*args, **kwargs) except Exception as e: # Convert exception to string @@ -148,15 +151,18 @@ def retry_on_exception(func): # If the specific error "context_length_exceeded" is present, simply return without retry if "context_length_exceeded" in err_str: + spinner_stop(spinner) raise Exception("context_length_exceeded") if "rate_limit_exceeded" in err_str: # Extracting the duration from the error string match = re.search(r"Please try again in (\d+)ms.", err_str) if match: + spinner = spinner_start(colored("Rate limited. Waiting...", 'yellow')) wait_duration = int(match.group(1)) / 1000 time.sleep(wait_duration) continue + spinner_stop(spinner) print(colored(f'There was a problem with request to openai API:', 'red')) print(err_str) diff --git a/pilot/utils/spinner.py b/pilot/utils/spinner.py index ba9e881..f5241ab 100644 --- a/pilot/utils/spinner.py +++ b/pilot/utils/spinner.py @@ -9,4 +9,5 @@ def spinner_start(text="Processing..."): def spinner_stop(spinner): - spinner.stop() + if spinner is not None: + spinner.stop()