diff --git a/src/openai_pygenerator/openai_pygenerator.py b/src/openai_pygenerator/openai_pygenerator.py index b8c45c9..e52d622 100644 --- a/src/openai_pygenerator/openai_pygenerator.py +++ b/src/openai_pygenerator/openai_pygenerator.py @@ -6,6 +6,7 @@ import openai from openai.error import APIError, RateLimitError, ServiceUnavailableError +from urllib3.exceptions import ReadTimeoutError Completion = Dict[str, str] Seconds = NewType("Seconds", int) @@ -88,7 +89,7 @@ def generate_completions( logger.debug("response = %s", result) for choice in result.choices: yield choice.message - except (RateLimitError, APIError, ServiceUnavailableError) as err: + except (ReadTimeoutError, RateLimitError, APIError, ServiceUnavailableError) as err: if isinstance(err, APIError) and not (err.http_status in [524, 502]): raise logger.warning("Error returned from openai API: %s", err)