Added retry functionality for allele variant determination

src/autobigs/engine/analysis/bigsdb.py

@@ -43,7 +43,8 @@ class BIGSdbMLSTProfiler(AbstractAsyncContextManager):
 
 class RemoteBIGSdbMLSTProfiler(BIGSdbMLSTProfiler):
 
-    def __init__(self, database_api: str, database_name: str, scheme_id: int):
+    def __init__(self, database_api: str, database_name: str, scheme_id: int, retry_requests: int = 5):
+        self._retry_limit = retry_requests
         self._database_name = database_name
         self._scheme_id = scheme_id
         self._base_url = f"{database_api}/db/{self._database_name}/schemes/{self._scheme_id}/"
@@ -57,11 +58,18 @@ class RemoteBIGSdbMLSTProfiler(BIGSdbMLSTProfiler):
         uri_path = "sequence"
         if isinstance(query_sequence_strings, str) or isinstance(query_sequence_strings, NamedString):
             query_sequence_strings = [query_sequence_strings]
+        
         for sequence_string in query_sequence_strings:
-            async with self._http_client.post(uri_path, json={
+            attempts = 0
+            success = False
+            last_error = None
+            while not success and attempts < self._retry_limit:
+                request = self._http_client.post(uri_path, json={
                     "sequence": sequence_string if isinstance(sequence_string, str) else sequence_string.sequence,
                     "partial_matches": True
-            }) as response:
+                })
+                try:
+                    async with request as response:
                         sequence_response: dict = await response.json()
 
                         if "exact_matches" in sequence_response:
@@ -91,6 +99,14 @@ class RemoteBIGSdbMLSTProfiler(BIGSdbMLSTProfiler):
                                 yield result_allele if isinstance(sequence_string, str) else (sequence_string.name, result_allele)
                         else:
                             raise NoBIGSdbMatchesException(self._database_name, self._scheme_id, sequence_string.name if isinstance(sequence_string, NamedString) else None)
+                except ConnectionResetError as e:
+                    last_error = e
+                    success = False
+                    await asyncio.sleep(5) # In case the connection issue is due to rate issues
+                else:
+                    success = True
+            if not success and last_error is not None:
+                raise last_error
             
     async def determine_mlst_st(self, alleles: Union[AsyncIterable[Union[Allele, tuple[str, Allele]]], Iterable[Union[Allele, tuple[str, Allele]]]]) -> Union[MLSTProfile, NamedMLSTProfile]:
         uri_path = "designations"
@@ -113,6 +129,7 @@ class RemoteBIGSdbMLSTProfiler(BIGSdbMLSTProfiler):
         request_json = {
             "designations": allele_request_dict
         }
+        
         async with self._http_client.post(uri_path, json=request_json) as response:
             response_json: dict = await response.json()
             allele_set: Set[Allele] = set()