retry on ReadTimeoutError

phelps-sg · Jun 29, 2023 · 429fb6f · 429fb6f
1 parent b7689fa
commit 429fb6f
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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)