From 44509cd6a1db8deefa9dfb9dfbeb6a058c03cc7c Mon Sep 17 00:00:00 2001 From: Sebastian Date: Fri, 17 Feb 2023 13:17:44 +0100 Subject: [PATCH] feat: Add OpenAIError to retry mechanism (#4178) * Add OpenAIError to retry mechanism. Use env variable for timeout for OpenAI request in PromptNode. * Updated retry in OpenAI embedding encoder as well. * Empty commit --- haystack/nodes/answer_generator/openai.py | 4 +++- haystack/nodes/prompt/prompt_node.py | 17 ++++++++++++++--- haystack/nodes/retriever/_openai_encoder.py | 4 +++- 3 files changed, 20 insertions(+), 5 deletions(-) diff --git a/haystack/nodes/answer_generator/openai.py b/haystack/nodes/answer_generator/openai.py index 003812a914..6cad709940 100644 --- a/haystack/nodes/answer_generator/openai.py +++ b/haystack/nodes/answer_generator/openai.py @@ -186,7 +186,9 @@ def __init__( logger.debug("Using GPT2TokenizerFast") self._hf_tokenizer: PreTrainedTokenizerFast = GPT2TokenizerFast.from_pretrained(tokenizer) - @retry_with_exponential_backoff(backoff_in_seconds=OPENAI_BACKOFF, max_retries=OPENAI_MAX_RETRIES) + @retry_with_exponential_backoff( + backoff_in_seconds=OPENAI_BACKOFF, max_retries=OPENAI_MAX_RETRIES, errors=(OpenAIRateLimitError, OpenAIError) + ) def predict( self, query: str, diff --git a/haystack/nodes/prompt/prompt_node.py b/haystack/nodes/prompt/prompt_node.py index 3cff64f1fd..8964585fac 100644 --- a/haystack/nodes/prompt/prompt_node.py +++ b/haystack/nodes/prompt/prompt_node.py @@ -20,7 +20,11 @@ from transformers.models.auto.modeling_auto import MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES from haystack import MultiLabel -from haystack.environment import HAYSTACK_REMOTE_API_BACKOFF_SEC, HAYSTACK_REMOTE_API_MAX_RETRIES +from haystack.environment import ( + HAYSTACK_REMOTE_API_BACKOFF_SEC, + HAYSTACK_REMOTE_API_MAX_RETRIES, + HAYSTACK_REMOTE_API_TIMEOUT_SEC, +) from haystack.errors import OpenAIError, OpenAIRateLimitError from haystack.modeling.utils import initialize_device_settings from haystack.nodes.base import BaseComponent @@ -435,8 +439,9 @@ def __init__( } @retry_with_exponential_backoff( - backoff_in_seconds=int(os.environ.get(HAYSTACK_REMOTE_API_BACKOFF_SEC, 5)), + backoff_in_seconds=float(os.environ.get(HAYSTACK_REMOTE_API_BACKOFF_SEC, 5)), max_retries=int(os.environ.get(HAYSTACK_REMOTE_API_MAX_RETRIES, 5)), + errors=(OpenAIRateLimitError, OpenAIError), ) def invoke(self, *args, **kwargs): """ @@ -478,7 +483,13 @@ def invoke(self, *args, **kwargs): "logit_bias": kwargs_with_defaults.get("logit_bias", {}), } headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"} - response = requests.request("POST", self.url, headers=headers, data=json.dumps(payload), timeout=30) + response = requests.request( + "POST", + self.url, + headers=headers, + data=json.dumps(payload), + timeout=float(os.environ.get(HAYSTACK_REMOTE_API_TIMEOUT_SEC, 30)), + ) res = json.loads(response.text) if response.status_code != 200: diff --git a/haystack/nodes/retriever/_openai_encoder.py b/haystack/nodes/retriever/_openai_encoder.py index 09134898a1..4ba2f0478d 100644 --- a/haystack/nodes/retriever/_openai_encoder.py +++ b/haystack/nodes/retriever/_openai_encoder.py @@ -104,7 +104,9 @@ def _ensure_text_limit(self, text: str) -> str: return decoded_string - @retry_with_exponential_backoff(backoff_in_seconds=OPENAI_BACKOFF, max_retries=OPENAI_MAX_RETRIES) + @retry_with_exponential_backoff( + backoff_in_seconds=OPENAI_BACKOFF, max_retries=OPENAI_MAX_RETRIES, errors=(OpenAIRateLimitError, OpenAIError) + ) def embed(self, model: str, text: List[str]) -> np.ndarray: payload = {"model": model, "input": text} headers = {"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"}