Merge pull request #875 from superannotateai/FRIDAY-5282

add integration for MM in upload_annotations

Author: nareksa

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

@@ -3702,6 +3702,7 @@ def upload_annotations(
         keep_status: bool | None = None,
         *,
         data_spec: Literal["default", "multimodal"] = "default",
+        integration: str | None = None,
     ):
         """Uploads a list of annotation dictionaries to the specified SuperAnnotate project or folder.
 
@@ -3723,6 +3724,10 @@ def upload_annotations(
                                      compact and modality-specific data representation.
         :type data_spec: str, optional
 
+        :param integration: The name of an existing integration on the SuperAnnotate platform, used to access external URLs in the annotations. Only supported for
+            Multimodal projects and data_spec="multimodal" and only applies to items being newly created — it has no effect on existing items.
+        :type integration: str, optional
+
         :return: A dictionary containing the results of the upload, categorized into successfully uploaded,
             failed, and skipped annotations.
         :rtype: dict
@@ -3759,6 +3764,17 @@ def upload_annotations(
                 keep_status=True,
                 data_spec='multimodal'
             )
+
+        Example Usage with private URLs signed via an integration::
+
+            # Upload annotations with private URLs using integration
+            sa_client.upload_annotations(
+                project="project1/folder1",
+                annotations=annotations,
+                keep_status=True,
+                data_spec="multimodal",
+                integration="AWS Main Bucket"
+            )
         """
         if keep_status is not None:
             warnings.warn(
@@ -3768,13 +3784,26 @@ def upload_annotations(
                 )
             )
         project, folder = self.controller.get_project_folder(project)
+        integration_entity = None
+        if integration:
+            if data_spec != "multimodal" or project.type != ProjectType.MULTIMODAL:
+                raise AppException(
+                    "Integration is only supported for Multimodal projects"
+                )
+            for i in self.controller.integrations.list().data:
+                if i.name == integration:
+                    integration_entity = i
+                    break
+            else:
+                raise AppException("Integration not found")
         response = self.controller.annotations.upload_multiple(
             project=project,
             folder=folder,
             annotations=annotations,
             keep_status=keep_status,
             user=self.controller.current_user,
             output_format=data_spec,
+            integration=integration_entity,
         )
         if response.errors:
             raise AppException(response.errors)

src/superannotate/lib/core/usecases/annotations.py

@@ -30,6 +30,7 @@
 from lib.core.entities import ConfigEntity
 from lib.core.entities import FolderEntity
 from lib.core.entities import ImageEntity
+from lib.core.entities import IntegrationEntity
 from lib.core.entities import ProjectEntity
 from lib.core.entities import UserEntity
 from lib.core.exceptions import AppException
@@ -1745,6 +1746,7 @@ def __init__(
         user: UserEntity,
         keep_status: bool = False,
         transform_version: str = None,
+        integration: IntegrationEntity = None,
     ):
         super().__init__(reporter)
         self._project = project
@@ -1758,6 +1760,7 @@ def __init__(
         self._transform_version = (
             "llmJsonV2" if transform_version is None else transform_version
         )
+        self._integration = integration
         self._category_name_to_id_map = {}
 
     @property
@@ -1895,7 +1898,12 @@ def attach_items(
             project=self._project,
             folder=folder,
             attachments=[
-                AttachmentEntity(name=item_name, url="") for item_name in item_names
+                AttachmentEntity(
+                    name=item_name,
+                    url="custom_llm",  # hardcoded for multimodal items
+                    integration_id=self._integration.id if self._integration else None,
+                )
+                for item_name in item_names
             ],
             service_provider=self._service_provider,
         ).execute()

src/superannotate/lib/infrastructure/controller.py

@@ -1432,6 +1432,7 @@ def upload_multiple(
         keep_status: bool,
         user: UserEntity,
         output_format: str = None,
+        integration: IntegrationEntity = None,
     ):
         if project.type == ProjectType.MULTIMODAL and output_format == "multimodal":
             use_case = usecases.UploadMultiModalAnnotationsUseCase(
@@ -1443,6 +1444,7 @@ def upload_multiple(
                 keep_status=keep_status,
                 user=user,
                 transform_version="llmJsonV2",
+                integration=integration,
             )
         else:
             use_case = usecases.UploadAnnotationsUseCase(

tests/integration/annotations/test_upload_annotations.py

@@ -1,3 +1,4 @@
+import base64
 import json
 import os
 import tempfile
@@ -257,6 +258,71 @@ def test_upload_with_integer_names(self):
                     f"{self.PROJECT_NAME}/test_folder", data_spec="multimodal"
                 )
 
+    def test_integration_not_found(self):
+        with open(self.JSONL_ANNOTATIONS_PATH) as f:
+            data = [json.loads(line) for line in f]
+        with self.assertRaisesRegex(AppException, "Integration not found"):
+            sa.upload_annotations(
+                self.PROJECT_NAME,
+                annotations=data,
+                data_spec="multimodal",
+                integration="non-existing-integration-xyz",
+            )
+
+    def test_integration_only_supported_for_multimodal_data_spec(self):
+        with open(self.JSONL_ANNOTATIONS_PATH) as f:
+            data = [json.loads(line) for line in f]
+        with self.assertRaisesRegex(
+            AppException, "Integration is only supported for Multimodal projects"
+        ):
+            sa.upload_annotations(
+                self.PROJECT_NAME,
+                annotations=data,
+                data_spec="default",
+                integration="any-integration",
+            )
+
+    def test_upload_with_existing_integration(self):
+        integrations = sa.get_integrations()
+        if not integrations:
+            self.skipTest("No integrations available in the team.")
+        integration = integrations[0]
+        with open(self.JSONL_ANNOTATIONS_PATH) as f:
+            data = [json.loads(line) for line in f]
+        response = sa.upload_annotations(
+            self.PROJECT_NAME,
+            annotations=data,
+            data_spec="multimodal",
+            integration=integration["name"],
+        )
+        assert len(response["succeeded"]) == 3
+
+        # Newly created items must carry the integration id used to sign URLs.
+        # The item metadata's integration_id isn't exposed by the SDK entities,
+        # so query the backend directly for it.
+        from lib.core.jsx_conditions import EmptyQuery
+        from lib.core.jsx_conditions import Join
+
+        project, folder = sa.controller.get_project_folder(
+            f"{self.PROJECT_NAME}/test_folder"
+        )
+        item_service = sa.controller.service_provider.item_service
+        client = item_service.client
+        entity_context = base64.b64encode(
+            f'{{"team_id":{client.team_id},"project_id":{project.id},'
+            f'"folder_id":{folder.id}}}'.encode()
+        ).decode()
+        raw_items = client.jsx_paginate(
+            url=item_service.URL_LIST,
+            chunk_size=2000,
+            body_query=EmptyQuery() & Join("metadata", ["path", "integration_id"]),
+            method="post",
+            headers={"x-sa-entity-context": entity_context},
+        ).data
+        assert len(raw_items) == 3
+        for item in raw_items:
+            assert item["metadata"]["integration_id"] == integration["id"]
+
     def test_download_annotations(self):
         with open(self.JSONL_ANNOTATIONS_PATH) as f:
             data = [json.loads(line) for line in f]