Skip to content
This repository has been archived by the owner on Nov 3, 2023. It is now read-only.

Commit

Permalink
Search query API timeout and retry (#4631)
Browse files Browse the repository at this point in the history
* The timeout added to the request

* retries

* removed unused const
  • Loading branch information
mojtaba-komeili authored Jun 27, 2022
1 parent 556048d commit 6af3a27
Showing 1 changed file with 23 additions and 6 deletions.
29 changes: 23 additions & 6 deletions parlai/agents/rag/retrieve_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -81,17 +81,34 @@ class SearchEngineRetriever(RetrieverAPI):
def __init__(self, opt: Opt):
super().__init__(opt=opt)
self.server_address = self._validate_server(opt.get('search_server'))
self._server_timeout = (
opt['search_server_timeout']
if opt.get('search_server_timeout', 0) > 0
else None
)
self._max_num_retries = opt.get('max_num_retries', 0)

def _query_search_server(self, query_term, n):
server = self.server_address
req = {'q': query_term, 'n': n}
logging.debug(f'sending search request to {server}')
server_response = requests.post(server, data=req)
resp_status = server_response.status_code
if resp_status == 200:
return server_response.json().get('response', None)
trials = []
while True:
try:
logging.debug(f'sending search request to {server}')
server_response = requests.post(
server, data=req, timeout=self._server_timeout
)
resp_status = server_response.status_code
trials.append(f'Response code: {resp_status}')
if resp_status == 200:
return server_response.json().get('response', None)
except requests.exceptions.Timeout:
if len(trials) > self._max_num_retries:
break
trials.append(f'Timeout after {self._server_timeout} seconds.')
logging.error(
f'Failed to retrieve data from server! Search server returned status {resp_status}'
f'Failed to retrieve data from server after {len(trials)+1} trials.'
f'\nFailed responses: {trials}'
)

def _validate_server(self, address):
Expand Down

0 comments on commit 6af3a27

Please sign in to comment.