Merge pull request #751 from superannotateai/get_annotations

update in get_annotations logic

Author: nareksa

src/superannotate/lib/app/interface/sdk_interface.py

@@ -501,7 +501,7 @@ def get_component_config(self, project: Union[NotEmptyStr, int], component_id: s
         """
 
         def retrieve_context(
-                component_data: List[dict], component_pk: str
+            component_data: List[dict], component_pk: str
         ) -> Tuple[bool, typing.Any]:
             try:
                 for component in component_data:
@@ -512,9 +512,9 @@ def retrieve_context(
                         if found:
                             return found, val
                     if (
-                            "id" in component and
-                            component["id"] == component_pk
-                            and component["type"] == "webComponent"
+                        "id" in component
+                        and component["id"] == component_pk
+                        and component["type"] == "webComponent"
                     ):
                         return True, json.loads(component.get("context"))
 

src/superannotate/lib/core/serviceproviders.py

@@ -490,6 +490,7 @@ def get_upload_chunks(
         self,
         project: entities.ProjectEntity,
         item_ids: List[int],
+        chunk_size: int = 1000,
     ) -> Dict[str, List]:
         raise NotImplementedError
 

src/superannotate/lib/infrastructure/services/annotation.py

@@ -18,6 +18,8 @@
 from lib.core.service_types import UploadAnnotationsResponse
 from lib.core.serviceproviders import BaseAnnotationService
 from lib.infrastructure.stream_data_handler import StreamedAnnotations
+from lib.infrastructure.utils import annotation_is_valid
+from lib.infrastructure.utils import divide_to_chunks
 
 try:
     from pydantic.v1 import parse_obj_as
@@ -170,21 +172,29 @@ def get_upload_chunks(
         self,
         project: entities.ProjectEntity,
         item_ids: List[int],
+        chunk_size: int = 1000,
     ) -> Dict[str, List]:
-        response_data = {"small": [], "large": []}
-        response = self.client.request(
-            url=urljoin(self.get_assets_provider_url(), self.URL_CLASSIFY_ITEM_SIZE),
-            method="POST",
-            params={"limit": len(item_ids)},
-            data={"project_id": project.id, "item_ids": item_ids},
-        )
-        if not response.ok:
-            raise AppException(response.error)
-        response_data["small"] = [
-            i["data"] for i in response.data.get("small", {}).values()
-        ]
-        response_data["large"] = response.data.get("large", [])
-        return response_data
+        small = []
+        large = []
+
+        chunks = divide_to_chunks(item_ids, chunk_size)
+        for chunk in chunks:
+            response = self.client.request(
+                method="POST",
+                url=urljoin(
+                    self.get_assets_provider_url(), self.URL_CLASSIFY_ITEM_SIZE
+                ),
+                params={"limit": len(chunk)},
+                data={
+                    "project_id": project.id,
+                    "item_ids": chunk,
+                },
+            )
+            if not response.ok:
+                raise AppException(response.error)
+            small.extend([i["data"] for i in response.data.get("small", {}).values()])
+            large.extend(response.data.get("large", []))
+        return {"small": small, "large": large}
 
     async def download_big_annotation(
         self,
@@ -218,6 +228,14 @@ async def download_big_annotation(
         ) as session:
             start_response = await session.request("post", url, params=query_params)
             res = await start_response.json()
+            if start_response.status > 299 or not annotation_is_valid(res):
+                logger.debug(
+                    f"Failed to download large annotation; item_id [{item_id}];"
+                    f" response: {res}; http_status: {start_response.status}"
+                )
+                raise AppException(
+                    f"Failed to download large annotation, ID: {item_id}"
+                )
             Path(download_path).mkdir(exist_ok=True, parents=True)
 
             dest_path = Path(download_path) / (item_name + ".json")

src/superannotate/lib/infrastructure/stream_data_handler.py

@@ -6,8 +6,12 @@
 from typing import Callable
 
 import aiohttp
+from lib.core.exceptions import AppException
+from lib.core.exceptions import BackendError
 from lib.core.reporter import Reporter
 from lib.infrastructure.services.http_client import AIOHttpSession
+from lib.infrastructure.utils import annotation_is_valid
+from lib.infrastructure.utils import async_retry_on_generator
 
 _seconds = 2**10
 TIMEOUT = aiohttp.ClientTimeout(
@@ -42,6 +46,7 @@ def get_json(self, data: bytes):
             self._reporter.log_error(f"Invalud chunk: {str(e)}")
             return None
 
+    @async_retry_on_generator((BackendError,))
     async def fetch(
         self,
         method: str,
@@ -59,6 +64,7 @@ async def fetch(
         buffer = ""
         line_groups = b""
         decoder = json.JSONDecoder()
+        data_received = False
         async for line in response.content.iter_any():
             line_groups += line
             try:
@@ -71,6 +77,19 @@ async def fetch(
                     if buffer.startswith(self.DELIMITER):
                         buffer = buffer[self.DELIMITER_LEN :]
                     json_obj, index = decoder.raw_decode(buffer)
+                    if not annotation_is_valid(json_obj):
+                        logger.warning(
+                            f"Invalid JSON detected in small annotations stream process, json: {json_obj}."
+                        )
+                        if data_received:
+                            raise AppException(
+                                "Invalid JSON detected in small annotations stream process."
+                            )
+                        else:
+                            raise BackendError(
+                                "Invalid JSON detected at the start of the small annotations stream process."
+                            )
+                    data_received = True
                     yield json_obj
                     if len(buffer[index:]) >= self.DELIMITER_LEN:
                         buffer = buffer[index + self.DELIMITER_LEN :]

src/superannotate/lib/infrastructure/utils.py

@@ -1,12 +1,17 @@
+import asyncio
+import logging
 import time
 from abc import ABC
 from abc import abstractmethod
+from functools import wraps
 from itertools import islice
 from pathlib import Path
 from typing import Any
+from typing import Callable
 from typing import Dict
 from typing import Optional
 from typing import Tuple
+from typing import Type
 from typing import Union
 
 from lib.core.entities import ProjectEntity
@@ -16,6 +21,9 @@
 from lib.infrastructure.services.work_management import WorkManagementService
 
 
+logger = logging.getLogger("sa")
+
+
 def divide_to_chunks(it, size):
     it = iter(it)
     return iter(lambda: tuple(islice(it, size)), ())
@@ -44,6 +52,66 @@ def extract_project_folder(user_input: Union[str, dict]) -> Tuple[str, Optional[
     raise PathError("Invalid project path")
 
 
+def async_retry_on_generator(
+    exceptions: Tuple[Type[Exception]],
+    retries: int = 3,
+    delay: float = 0.3,
+    backoff: float = 0.3,
+):
+    """
+    An async retry decorator that retries a function only on specific exceptions.
+
+    Parameters:
+        exceptions (tuple): Tuple of exception classes to retry on.
+        retries (int): Number of retry attempts.
+        delay (float): Initial delay between retries in seconds.
+        backoff (float): Factor to increase the delay after each failure.
+    """
+
+    def decorator(func: Callable):
+        @wraps(func)
+        async def wrapper(*args, **kwargs):
+            attempt = 0
+            current_delay = delay
+            raised_exception = None
+
+            while attempt < retries:
+                try:
+                    async for v in func(*args, **kwargs):
+                        yield v
+                    return
+                except exceptions as e:
+                    raised_exception = e
+                    logger.debug(
+                        f"Attempt {attempt + 1}/{retries} failed with error: {e}. "
+                        f"Retrying in {current_delay} seconds..."
+                    )
+                    await asyncio.sleep(current_delay)
+                    current_delay += backoff  # Exponential backoff
+                finally:
+                    attempt += 1
+            if raised_exception:
+                logger.error(
+                    f"All {retries} attempts failed due to {raised_exception}."
+                )
+                raise raised_exception
+
+        return wrapper
+
+    return decorator
+
+
+def annotation_is_valid(annotation: dict) -> bool:
+    annotation_keys = annotation.keys()
+    if (
+        "errors" in annotation_keys
+        or "error" in annotation_keys
+        or "metadata" not in annotation_keys
+    ):
+        return False
+    return True
+
+
 class BaseCachedWorkManagementRepository(ABC):
     def __init__(self, ttl_seconds: int, work_management: WorkManagementService):
         self.ttl_seconds = ttl_seconds

tests/integration/annotations/test_upload_annotations_from_folder_to_project.py

@@ -166,8 +166,3 @@ def test_annotation_folder_upload_download(self):
                 contents1 = f1.read()
                 contents2 = f2.read()
                 assert contents1 == contents2
-
-
-def test():
-    a = sa.get_user_metadata(pk=244700, include=["custom_fields"])
-    print(a)