Coverage for haystack/utils/requests_utils.py: 100%

1# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai> 

2# 

3# SPDX-License-Identifier: Apache-2.0 

4 

5import logging 

6from typing import Any 

7 

8import httpx 

9from tenacity import after_log, before_log, retry, retry_if_exception_type, stop_after_attempt, wait_exponential 

10 

11# NOTE: this uses the standard library logger (not `haystack.logging`) on purpose: tenacity's `before_log`/`after_log` 

12# call the logger with positional arguments, which Haystack's keyword-only patched logger would reject. We still name 

13# it with `__name__` so it lives under the `haystack` namespace and is picked up by `configure_logging`. 

14logger = logging.getLogger(__name__) 

15 

16 

17def request_with_retry( 

18 attempts: int = 3, status_codes_to_retry: list[int] | None = None, **kwargs: Any 

19) -> httpx.Response: 

20 """ 

21 Executes an HTTP request with a configurable exponential backoff retry on failures. 

22 

23 Usage example: 

24 <!-- test-ignore --> 

25 ```python 

26 from haystack.utils import request_with_retry 

27 

28 # Sending an HTTP request with default retry configs 

29 res = request_with_retry(method="GET", url="https://example.com") 

30 

31 # Sending an HTTP request with custom number of attempts 

32 res = request_with_retry(method="GET", url="https://example.com", attempts=10) 

33 

34 # Sending an HTTP request with custom HTTP codes to retry 

35 res = request_with_retry(method="GET", url="https://example.com", status_codes_to_retry=[408, 503]) 

36 

37 # Sending an HTTP request with custom timeout in seconds 

38 res = request_with_retry(method="GET", url="https://example.com", timeout=5) 

39 

40 # Sending an HTTP request with custom headers 

41 res = request_with_retry(method="GET", url="https://example.com", headers={"Authorization": "Bearer <token>"}) 

42 

43 # Sending a POST request 

44 res = request_with_retry(method="POST", url="https://example.com", json={"key": "value"}, attempts=10) 

45 

46 # Retry all 5xx status codes 

47 res = request_with_retry(method="GET", url="https://example.com", status_codes_to_retry=list(range(500, 600))) 

48 ``` 

49 

50 :param attempts: 

51 Maximum number of attempts to retry the request. 

52 :param status_codes_to_retry: 

53 List of HTTP status codes that will trigger a retry. 

54 When param is `None`, HTTP 408, 418, 429 and 503 will be retried. 

55 :param kwargs: 

56 Optional arguments that `httpx.Client.request` accepts. 

57 :returns: 

58 The `httpx.Response` object. 

59 """ 

60 

61 if status_codes_to_retry is None: 

62 status_codes_to_retry = [408, 418, 429, 503] 

63 

64 # Pop `timeout` once, before the retry loop. 

65 timeout = kwargs.pop("timeout", 10) 

66 

67 @retry( 

68 reraise=True, 

69 wait=wait_exponential(), 

70 retry=retry_if_exception_type((httpx.HTTPError, TimeoutError)), 

71 stop=stop_after_attempt(attempts), 

72 before=before_log(logger, logging.DEBUG), 

73 after=after_log(logger, logging.DEBUG), 

74 ) 

75 def run() -> httpx.Response: 

76 with httpx.Client() as client: 

77 res = client.request(**kwargs, timeout=timeout) 

78 

79 if res.status_code in status_codes_to_retry: 

80 # We raise only for the status codes that must trigger a retry 

81 res.raise_for_status() 

82 

83 return res 

84 

85 res = run() 

86 # We raise here too in case the request failed with a status code that 

87 # won't trigger a retry, this way the call will still cause an explicit exception 

88 res.raise_for_status() 

89 return res 

90 

91 

92async def async_request_with_retry( 

93 attempts: int = 3, status_codes_to_retry: list[int] | None = None, **kwargs: Any 

94) -> httpx.Response: 

95 """ 

96 Executes an asynchronous HTTP request with a configurable exponential backoff retry on failures. 

97 

98 Usage example: 

99 ```python 

100 import asyncio 

101 from haystack.utils import async_request_with_retry 

102 

103 # Sending an async HTTP request with default retry configs 

104 async def example(): 

105 res = await async_request_with_retry(method="GET", url="https://example.com") 

106 return res 

107 

108 # Sending an async HTTP request with custom number of attempts 

109 async def example_with_attempts(): 

110 res = await async_request_with_retry(method="GET", url="https://example.com", attempts=10) 

111 return res 

112 

113 # Sending an async HTTP request with custom HTTP codes to retry 

114 async def example_with_status_codes(): 

115 res = await async_request_with_retry(method="GET", url="https://example.com", status_codes_to_retry=[408, 503]) 

116 return res 

117 

118 # Sending an async HTTP request with custom timeout in seconds 

119 async def example_with_timeout(): 

120 res = await async_request_with_retry(method="GET", url="https://example.com", timeout=5) 

121 return res 

122 

123 # Sending an async HTTP request with custom headers 

124 async def example_with_headers(): 

125 headers = {"Authorization": "Bearer <my_token_here>"} 

126 res = await async_request_with_retry(method="GET", url="https://example.com", headers=headers) 

127 return res 

128 

129 # All of the above combined 

130 async def example_combined(): 

131 headers = {"Authorization": "Bearer <my_token_here>"} 

132 res = await async_request_with_retry( 

133 method="GET", 

134 url="https://example.com", 

135 headers=headers, 

136 attempts=10, 

137 status_codes_to_retry=[408, 503], 

138 timeout=5 

139 ) 

140 return res 

141 

142 # Sending an async POST request 

143 async def example_post(): 

144 res = await async_request_with_retry( 

145 method="POST", 

146 url="https://example.com", 

147 json={"key": "value"}, 

148 attempts=10 

149 ) 

150 return res 

151 

152 # Retry all 5xx status codes 

153 async def example_5xx(): 

154 res = await async_request_with_retry( 

155 method="GET", 

156 url="https://example.com", 

157 status_codes_to_retry=list(range(500, 600)) 

158 ) 

159 return res 

160 ``` 

161 

162 :param attempts: 

163 Maximum number of attempts to retry the request. 

164 :param status_codes_to_retry: 

165 List of HTTP status codes that will trigger a retry. 

166 When param is `None`, HTTP 408, 418, 429 and 503 will be retried. 

167 :param kwargs: 

168 Optional arguments that `httpx.AsyncClient.request` accepts. 

169 :returns: 

170 The `httpx.Response` object. 

171 """ 

172 

173 if status_codes_to_retry is None: 

174 status_codes_to_retry = [408, 418, 429, 503] 

175 

176 # Pop `timeout` once, before the retry loop. 

177 timeout = kwargs.pop("timeout", 10) 

178 

179 @retry( 

180 reraise=True, 

181 wait=wait_exponential(), 

182 retry=retry_if_exception_type((httpx.HTTPError, TimeoutError)), 

183 stop=stop_after_attempt(attempts), 

184 before=before_log(logger, logging.DEBUG), 

185 after=after_log(logger, logging.DEBUG), 

186 ) 

187 async def run() -> httpx.Response: 

188 async with httpx.AsyncClient() as client: 

189 res = await client.request(**kwargs, timeout=timeout) 

190 

191 if res.status_code in status_codes_to_retry: 

192 # We raise only for the status codes that must trigger a retry 

193 res.raise_for_status() 

194 

195 return res 

196 

197 res = await run() 

198 # We raise here too in case the request failed with a status code that 

199 # won't trigger a retry, this way the call will still cause an explicit exception 

200 res.raise_for_status() 

201 return res