feat: Allow users to set a timeout for remote APIs

haystack/environment.py

@@ -3,18 +3,19 @@
 import platform
 import sys
 from typing import Any, Dict
+
 import torch
 import transformers
 
 from haystack import __version__
 
-
 HAYSTACK_EXECUTION_CONTEXT = "HAYSTACK_EXECUTION_CONTEXT"
 HAYSTACK_DOCKER_CONTAINER = "HAYSTACK_DOCKER_CONTAINER"
 
 # Any remote API (OpenAI, Cohere etc.)
 HAYSTACK_REMOTE_API_BACKOFF_SEC = "HAYSTACK_REMOTE_API_BACKOFF_SEC"
 HAYSTACK_REMOTE_API_MAX_RETRIES = "HAYSTACK_REMOTE_API_MAX_RETRIES"
+HAYSTACK_REMOTE_API_TIMEOUT_SEC = "HAYSTACK_REMOTE_API_TIMEOUT_SEC"
 
 env_meta_data: Dict[str, Any] = {}
 

haystack/nodes/answer_generator/openai.py

@@ -1,12 +1,18 @@
 import json
 import logging
+import os
+import platform
 import sys
 from typing import List, Optional, Tuple, Union
-import platform
 
 import requests
 
 from haystack import Document
+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.nodes.answer_generator import BaseGenerator
 from haystack.utils.reflection import retry_with_exponential_backoff
@@ -29,6 +35,11 @@
     from transformers import GPT2TokenizerFast, PreTrainedTokenizerFast
 
 
+OPENAI_TIMEOUT = float(os.environ.get(HAYSTACK_REMOTE_API_TIMEOUT_SEC, 30))
+OPENAI_BACKOFF = float(os.environ.get(HAYSTACK_REMOTE_API_BACKOFF_SEC, 10))
+OPENAI_MAX_RETRIES = int(os.environ.get(HAYSTACK_REMOTE_API_MAX_RETRIES, 5))
+
+
 class OpenAIAnswerGenerator(BaseGenerator):
     """
     Uses the GPT-3 models from the OpenAI API to generate Answers based on the Documents it receives.
@@ -117,13 +128,13 @@ def __init__(
             logger.debug("Using GPT2TokenizerFast")
             self._hf_tokenizer: PreTrainedTokenizerFast = GPT2TokenizerFast.from_pretrained(tokenizer)
 
-    @retry_with_exponential_backoff(backoff_in_seconds=10, max_retries=5)
+    @retry_with_exponential_backoff(backoff_in_seconds=OPENAI_BACKOFF, max_retries=OPENAI_MAX_RETRIES)
     def predict(
         self,
         query: str,
         documents: List[Document],
         top_k: Optional[int] = None,
-        timeout: Union[float, Tuple[float, float]] = 10.0,
+        timeout: Union[float, Tuple[float, float]] = OPENAI_TIMEOUT,
     ):
         """
         Use the loaded QA model to generate Answers for a query based on the Documents it receives.

haystack/nodes/retriever/_embedding_encoder.py

@@ -1,5 +1,6 @@
 import json
 import logging
+import os
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
 
@@ -12,6 +13,11 @@
 from tqdm.auto import tqdm
 from transformers import AutoModel, AutoTokenizer
 
+from haystack.environment import (
+    HAYSTACK_REMOTE_API_BACKOFF_SEC,
+    HAYSTACK_REMOTE_API_MAX_RETRIES,
+    HAYSTACK_REMOTE_API_TIMEOUT_SEC,
+)
 from haystack.errors import CohereError
 from haystack.modeling.data_handler.dataloader import NamedDataLoader
 from haystack.modeling.data_handler.dataset import convert_features_to_dataset, flatten_rename
@@ -27,6 +33,11 @@
     from haystack.nodes.retriever import EmbeddingRetriever
 
 
+COHERE_TIMEOUT = float(os.environ.get(HAYSTACK_REMOTE_API_TIMEOUT_SEC, 30))
+COHERE_BACKOFF = float(os.environ.get(HAYSTACK_REMOTE_API_BACKOFF_SEC, 10))
+COHERE_MAX_RETRIES = int(os.environ.get(HAYSTACK_REMOTE_API_MAX_RETRIES, 5))
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -321,11 +332,13 @@ def __init__(self, retriever: "EmbeddingRetriever"):
             "multilingual-22-12",
         )
 
-    @retry_with_exponential_backoff(backoff_in_seconds=10, max_retries=5, errors=(CohereError,))
+    @retry_with_exponential_backoff(
+        backoff_in_seconds=COHERE_BACKOFF, max_retries=COHERE_MAX_RETRIES, errors=(CohereError,)
+    )
     def embed(self, model: str, text: List[str]) -> np.ndarray:
         payload = {"model": model, "texts": text, "truncate": "END"}
         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=COHERE_TIMEOUT)
         res = json.loads(response.text)
         if response.status_code != 200:
             raise CohereError(response.text, status_code=response.status_code)

haystack/nodes/retriever/_openai_encoder.py

@@ -1,20 +1,25 @@
 import json
 import logging
+import os
+import platform
 import sys
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
-import platform
 
 import numpy as np
 import requests
 from tqdm.auto import tqdm
 
+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.nodes.retriever._base_embedding_encoder import _BaseEmbeddingEncoder
 from haystack.schema import Document
 from haystack.utils.reflection import retry_with_exponential_backoff
 
-
 if TYPE_CHECKING:
     from haystack.nodes.retriever import EmbeddingRetriever
 
@@ -36,6 +41,11 @@
     from transformers import GPT2TokenizerFast, PreTrainedTokenizerFast
 
 
+OPENAI_TIMEOUT = float(os.environ.get(HAYSTACK_REMOTE_API_TIMEOUT_SEC, 30))
+OPENAI_BACKOFF = float(os.environ.get(HAYSTACK_REMOTE_API_BACKOFF_SEC, 10))
+OPENAI_MAX_RETRIES = int(os.environ.get(HAYSTACK_REMOTE_API_MAX_RETRIES, 5))
+
+
 logger = logging.getLogger(__name__)
 
 
@@ -94,11 +104,11 @@ def _ensure_text_limit(self, text: str) -> str:
 
         return decoded_string
 
-    @retry_with_exponential_backoff(backoff_in_seconds=10, max_retries=5)
+    @retry_with_exponential_backoff(backoff_in_seconds=OPENAI_BACKOFF, max_retries=OPENAI_MAX_RETRIES)
     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"}
-        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=OPENAI_TIMEOUT)
         res = json.loads(response.text)
 
         if response.status_code != 200: