From ab972ccf6bee199605c7c66db08b784ca595eb85 Mon Sep 17 00:00:00 2001 From: Ximi Hoque Date: Wed, 5 Nov 2025 22:54:41 +0530 Subject: [PATCH 01/13] S3 connection crud refactored, schemas refactored --- driver.py | 55 +++---- drivers/connectors.py | 42 +++++ labellerr/core/connectors/__init__.py | 110 +++++++------- labellerr/core/connectors/connections.py | 143 ++++++------------ labellerr/core/connectors/s3_connection.py | 62 ++++++-- labellerr/core/constants.py | 3 +- labellerr/core/schemas/__init__.py | 6 +- .../schemas/{connections.py => connectors.py} | 6 + 8 files changed, 224 insertions(+), 203 deletions(-) create mode 100644 drivers/connectors.py rename labellerr/core/schemas/{connections.py => connectors.py} (95%) diff --git a/driver.py b/driver.py index dbf9b61..45efd98 100644 --- a/driver.py +++ b/driver.py @@ -5,7 +5,7 @@ from labellerr.client import LabellerrClient from labellerr.core.datasets import LabellerrDataset, create_dataset -from labellerr.core.connectors import LabellerrS3Connection +from labellerr.core.connectors import LabellerrS3Connection, LabellerrConnection from labellerr.core.files import LabellerrFile from labellerr.core.projects import ( LabellerrProject, @@ -25,26 +25,7 @@ ) from labellerr.core.autolabel import LabellerrAutoLabel -# Set logging level to DEBUG -logging.basicConfig(level=logging.DEBUG) -load_dotenv() - -API_KEY = os.getenv("API_KEY") -API_SECRET = os.getenv("API_SECRET") -CLIENT_ID = os.getenv("CLIENT_ID") - -if not all([API_KEY, API_SECRET, CLIENT_ID]): - raise ValueError( - "API_KEY, API_SECRET, and CLIENT_ID must be set in environment variables" - ) - -# Initialize client -client = LabellerrClient( - api_key=API_KEY, - api_secret=API_SECRET, - client_id=CLIENT_ID, -) # if os.getenv("CREATE_DATASET", "").lower() == "true": # from labellerr import schemas @@ -160,25 +141,25 @@ # name="Amazon S3 Export Test", # description="Amazon S3 Export Test", # ))) -project = LabellerrProject(client=client, project_id="") - -export = project.create_export( - export_config=CreateExportParams( - export_name="Amazon S3 Export Test", - export_description="Amazon S3 Export Test", - export_format="json", - statuses=["review"], - connection_id=os.getenv("AWS_EXPORT_CONNECTION_ID"), - export_destination="s3", - export_folder_path="", # pattern - bucket_name/path/to/folder/ - the last slash is important - ) -) +# project = LabellerrProject(client=client, project_id="") + +# export = project.create_export( +# export_config=CreateExportParams( +# export_name="Amazon S3 Export Test", +# export_description="Amazon S3 Export Test", +# export_format="json", +# statuses=["review"], +# connection_id=os.getenv("AWS_EXPORT_CONNECTION_ID"), +# export_destination="s3", +# export_folder_path="", # pattern - bucket_name/path/to/folder/ - the last slash is important +# ) +# ) -print(f"Export created: {export.report_id}") -print(f"Current status: {export._status}") +# print(f"Export created: {export.report_id}") +# print(f"Current status: {export._status}") # Uncomment to poll until completion: -final_status = export.status() -print(f"Final status: {final_status}") +# final_status = export.status() +# print(f"Final status: {final_status}") # print(autolabel.train(training_request=TrainingRequest(model_id="yolov11", hyperparameters=Hyperparameters(epochs=10), slice_id='', min_samples_per_class=100, job_name="Yolo V11 Training"))) # dataset = LabellerrDataset(client=client, dataset_id="") diff --git a/drivers/connectors.py b/drivers/connectors.py new file mode 100644 index 0000000..f9f114c --- /dev/null +++ b/drivers/connectors.py @@ -0,0 +1,42 @@ +import logging +import os +from dotenv import load_dotenv +from labellerr.client import LabellerrClient +from labellerr.core.connectors import LabellerrConnection, list_connections, delete_connection, create_connection, AWSConnectionParams +from labellerr.core.schemas import ConnectionType, ConnectorType + +# Set logging level to DEBUG +logging.basicConfig(level=logging.DEBUG) + +load_dotenv() + +API_KEY = os.getenv("API_KEY") +API_SECRET = os.getenv("API_SECRET") +CLIENT_ID = os.getenv("CLIENT_ID") + +if not all([API_KEY, API_SECRET, CLIENT_ID]): + raise ValueError( + "API_KEY, API_SECRET, and CLIENT_ID must be set in environment variables" + ) + +# Initialize client +client = LabellerrClient( + api_key=API_KEY, + api_secret=API_SECRET, + client_id=CLIENT_ID, +) + +# connection = create_connection(client=client, connector_type=ConnectorType._S3, params=AWSConnectionParams( +# aws_access_key=os.getenv("AWS_KEY"), +# aws_secrets_key=os.getenv("AWS_SECRET"), +# s3_path="s3://amazon-s3-sync-test/labellerr-processed/videos/", # this path is not part of the connection but needed to test the connection on the desired path. +# # This can be dynamically changed while using the connection for creating datasets. +# connection_type=ConnectionType._IMPORT, +# name="Amazon S3 Import Test", +# description="Amazon S3 Import Test", +# )) +# print(f"Connection created: {connection.connection_id}") +# connection = LabellerrConnection(client=client, connection_id='2a30c044-57f7-42c7-8290-bab3bbac0ebc') +# print('connection type', connection.connection_type) + +# print (connection.test(s3_path="s3://amazon-s3-sync-test/labellerr-processed/videos/datasets/", connection_type=ConnectionType._IMPORT)) \ No newline at end of file diff --git a/labellerr/core/connectors/__init__.py b/labellerr/core/connectors/__init__.py index cda7749..688bb91 100644 --- a/labellerr/core/connectors/__init__.py +++ b/labellerr/core/connectors/__init__.py @@ -1,6 +1,8 @@ +import uuid from typing import TYPE_CHECKING -from ...schemas import AWSConnectionParams +from .. import constants +from ...schemas import AWSConnectionParams, ConnectionType, ConnectorType from .connections import LabellerrConnection from .gcs_connection import GCSConnection as LabellerrGCSConnection from .s3_connection import S3Connection as LabellerrS3Connection @@ -13,60 +15,66 @@ def create_connection( client: "LabellerrClient", - connector_type: str, - client_id: str, - connector_config: dict, -): - """ - Sets up cloud connector (GCP/AWS) for dataset creation using factory pattern. + connector_type: ConnectorType, + params: AWSConnectionParams, +) -> LabellerrS3Connection | LabellerrGCSConnection: + if connector_type == ConnectorType._S3: + return LabellerrS3Connection.create_connection(client, params) + elif connector_type == ConnectorType._GCS: + return LabellerrGCSConnection.create_connection(client, params) + else: + raise ValueError(f"Unsupported connector type: {connector_type}") + +def list_connections( + client: "LabellerrClient", + connector: ConnectorType, + connection_type: ConnectionType = None, + ) -> list[LabellerrGCSConnection | LabellerrS3Connection]: + """ + Lists connections for a client :param client: LabellerrClient instance - :param connector_type: Type of connector ('gcp' or 'aws') - :param client_id: Client ID - :param connector_config: Configuration dictionary for the connector - :return: Connection ID (str) for quick connection or full response (dict) for full connection + :param connection_type: Type of connection (import/export) + :param connector: Optional connector type filter (s3, gcs, etc.) + :return: List of connections """ - import logging - from ..exceptions import InvalidConnectionError + unique_id = str(uuid.uuid4()) + url = f"{constants.BASE_URL}/connectors/connections/list" + + params = { + "client_id": client.client_id, + "uuid": unique_id, + "connector": connector, + } + if connection_type: + params["connection_type"] = connection_type + extra_headers = {"email_id": client.api_key} - try: - if connector_type == "gcs": - from .gcs_connection import GCSConnection + response = client.make_request( + "GET", + url, + extra_headers=extra_headers, + request_id=unique_id, + params=params + ) + return [LabellerrConnection(client, connection["connection_id"]) for connection in response.get("response", [])] + +def delete_connection(client: "LabellerrClient", connection_id: str): + """ + Deletes a connector connection by ID. + :param connection_id: The ID of the connection to delete + :return: Parsed JSON response + """ + request_id = str(uuid.uuid4()) + url = ( + f"{constants.BASE_URL}/connectors/connections/delete" + f"?client_id={client.client_id}&uuid={request_id}" + ) - return GCSConnection.create_connection(client, connector_config) - elif connector_type == "s3": - from .s3_connection import S3Connection + extra_headers = {"email_id": client.api_key} - # Determine which method to call based on config parameters - # Full connection has: aws_access_key, aws_secrets_key, s3_path, name, description - # Quick connection has: bucket_name, folder_path, access_key_id, secret_access_key - if "aws_access_key" in connector_config and "name" in connector_config: - # Full connection flow - creates a saved connection - return S3Connection.setup_full_connection( - client, - AWSConnectionParams( - client_id=connector_config["client_id"], - aws_access_key=connector_config["aws_access_key"], - aws_secrets_key=connector_config["aws_secrets_key"], - s3_path=connector_config["s3_path"], - data_type=connector_config["data_type"], - name=connector_config["name"], - description=connector_config["description"], - connection_type=connector_config.get( - "connection_type", "import" - ), - ), - ) - else: - # Quick connection flow - for dataset creation - return S3Connection.create_connection( - client, client_id, connector_config - ) - else: - raise InvalidConnectionError( - f"Unsupported connector type: {connector_type}" - ) - except Exception as e: - logging.error(f"Failed to setup {connector_type} connector: {e}") - raise + response = client.make_request( + "POST", url, extra_headers=extra_headers, request_id=request_id, json={"connection_id": connection_id} + ) + return response.get("response", None) diff --git a/labellerr/core/connectors/connections.py b/labellerr/core/connectors/connections.py index b18a6d4..718ea22 100644 --- a/labellerr/core/connectors/connections.py +++ b/labellerr/core/connectors/connections.py @@ -3,8 +3,9 @@ import uuid from abc import ABCMeta, abstractmethod from typing import TYPE_CHECKING, Dict - +import logging from .. import client_utils, constants +from ..schemas import ConnectionType, ConnectorType from ..exceptions import InvalidConnectionError, InvalidDatasetIDError if TYPE_CHECKING: @@ -23,24 +24,26 @@ def _register(cls, connection_type, connection_class): @staticmethod def get_connection(client: "LabellerrClient", connection_id: str): """Get connection from Labellerr API""" - # ------------------------------- [needs refactoring after we consolidate api_calls into one function ] --------------------------------- - unique_id = str(uuid.uuid4()) - url = ( - f"{constants.BASE_URL}/connections/{connection_id}?client_id={client.client_id}" - f"&uuid={unique_id}" - ) - headers = client_utils.build_headers( - api_key=client.api_key, - api_secret=client.api_secret, - client_id=client.client_id, - extra_headers={"content-type": "application/json"}, - ) - response = client_utils.request( - "GET", url, headers=headers, request_id=unique_id + assert connection_id, "Connection ID is can't be empty" + unique_id = str(uuid.uuid4()) + url = f"{constants.BASE_URL}/connectors/connections/{connection_id}/details" + + params = { + "client_id": client.client_id, + "uuid": unique_id + } + + extra_headers = {"content-type": "application/json"} + + response = client.make_request( + "GET", + url, + extra_headers=extra_headers, + request_id=unique_id, + params=params ) return response.get("response", None) - # ------------------------------- [needs refactoring after we consolidate api_calls into one function ] --------------------------------- """Metaclass that combines ABC functionality with factory pattern""" @@ -55,15 +58,10 @@ def __call__(cls, client, connection_id, **kwargs): connection_data = cls.get_connection(client, connection_id) if connection_data is None: raise InvalidDatasetIDError(f"Connection not found: {connection_id}") - connection_type = connection_data.get("connection_type") - if connection_type not in constants.CONNECTION_TYPES: - raise InvalidConnectionError( - f"Connection type not supported: {connection_type}" - ) - - connection_class = cls._registry.get(connection_type) + connector = connection_data.get("connector") + connection_class = cls._registry.get(connector) if connection_class is None: - raise InvalidConnectionError(f"Unknown connection type: {connection_type}") + raise InvalidConnectionError(f"Unknown connector type: {connector}") kwargs["connection_data"] = connection_data return connection_class(client, connection_id, **kwargs) @@ -74,84 +72,37 @@ class LabellerrConnection(metaclass=LabellerrConnectionMeta): def __init__(self, client: "LabellerrClient", connection_id: str, **kwargs): self.client = client self._connection_id_input = connection_id - self.connection_data = kwargs["connection_data"] + self.__connection_data = kwargs["connection_data"] + + @property + def name(self): + return self.__connection_data.get("name") + + @property + def description(self): + return self.__connection_data.get("description") @property def connection_id(self): - return self.connection_data.get("connection_id") + return self.__connection_data.get("connection_id") @property def connection_type(self): - return self.connection_data.get("connection_type") + return self.__connection_data.get("connection_type") + + @property + def connector(self): + return self.__connection_data.get("connector") + + @property + def created_at(self): + return self.__connection_data.get("created_at") + + @property + def created_by(self): + return self.__connection_data.get("created_by") @abstractmethod - def test_connection(self): + def test(self): """Each connection type must implement its own connection testing logic""" - pass - - def list_connections( - self, - connection_type: str, - connector: str = None, - ) -> list: - """ - List connections for a client - :param connection_type: Type of connection (import/export) - :param connector: Optional connector type filter (s3, gcs, etc.) - :return: List of connections - """ - request_uuid = str(uuid.uuid4()) - list_connection_url = ( - f"{constants.BASE_URL}/connectors/connections/list" - f"?client_id={self.client.client_id}&uuid={request_uuid}&connection_type={connection_type}" - ) - - if connector: - list_connection_url += f"&connector={connector}" - - headers = client_utils.build_headers( - api_key=self.client.api_key, - api_secret=self.client.api_secret, - client_id=self.client.client_id, - extra_headers={"email_id": self.client.api_key}, - ) - - return client_utils.request( - "GET", list_connection_url, headers=headers, request_id=request_uuid - ) - - def delete_connection(self, connection_id: str): - """ - Deletes a connector connection by ID. - :param connection_id: The ID of the connection to delete - :return: Parsed JSON response - """ - import json - - from ... import schemas - - # Validate parameters using Pydantic - params = schemas.DeleteConnectionParams( - client_id=self.client.client_id, connection_id=connection_id - ) - request_uuid = str(uuid.uuid4()) - delete_url = ( - f"{constants.BASE_URL}/connectors/connections/delete" - f"?client_id={params.client_id}&uuid={request_uuid}" - ) - - headers = client_utils.build_headers( - api_key=self.client.api_key, - api_secret=self.client.api_secret, - client_id=self.client.client_id, - extra_headers={ - "content-type": "application/json", - "email_id": self.client.api_key, - }, - ) - - payload = json.dumps({"connection_id": params.connection_id}) - - return client_utils.request( - "POST", delete_url, headers=headers, data=payload, request_id=request_uuid - ) + pass \ No newline at end of file diff --git a/labellerr/core/connectors/s3_connection.py b/labellerr/core/connectors/s3_connection.py index 0e8662c..d65095b 100644 --- a/labellerr/core/connectors/s3_connection.py +++ b/labellerr/core/connectors/s3_connection.py @@ -1,31 +1,26 @@ import json import uuid -from typing import TYPE_CHECKING - -from ..schemas import AWSConnectionParams, AWSConnectionTestParams +from ..client import LabellerrClient +from ..schemas import AWSConnectionParams, AWSConnectionTestParams, ConnectionType from .. import client_utils, constants from .connections import LabellerrConnection, LabellerrConnectionMeta - -if TYPE_CHECKING: - from labellerr import LabellerrClient - - +import logging class S3Connection(LabellerrConnection): @staticmethod def test_connection( client: "LabellerrClient", params: AWSConnectionTestParams - ) -> bool: + ) -> dict: """ Tests an AWS S3 connection. :param client: The LabellerrClient instance :param params: The AWS connection parameters :return: True if the connection is successful, False otherwise """ - request_uuid = str(uuid.uuid4()) + request_id = str(uuid.uuid4()) test_connection_url = ( f"{constants.BASE_URL}/connectors/connections/test" - f"?client_id={client.client_id}&uuid={request_uuid}" + f"?client_id={client.client_id}&uuid={request_id}" ) headers = client_utils.build_headers( @@ -55,14 +50,44 @@ def test_connection( k: v for k, v in headers.items() if k.lower() != "content-type" } - client_utils.request( + response = client_utils.request( "POST", test_connection_url, headers=headers_without_content_type, files=test_request, - request_id=request_uuid, + request_id=request_id, + ) + return response.get("response", {}) + + def test(self, s3_path: str, connection_type: ConnectionType): + request_id = str(uuid.uuid4()) + test_connection_url = ( + f"{constants.BASE_URL}/connectors/connections/test" + f"?client_id={self.client.client_id}&uuid={request_id}" + ) + + headers = client_utils.build_headers( + api_key=self.client.api_key, + api_secret=self.client.api_secret, + client_id=self.client.client_id, + extra_headers={"email_id": self.client.api_key}, + ) + + # Test endpoint also expects multipart/form-data format + test_request = { + "connector": (None, "s3"), + "path": (None, s3_path), + "connection_type": (None, connection_type.value), + "connection_id": (None, self.connection_id), + } + response = client_utils.request( + "POST", + test_connection_url, + headers=headers, + files=test_request, + request_id=request_id, ) - return True + return response.get("response", {}) @staticmethod def create_connection( @@ -76,6 +101,11 @@ def create_connection( :return: Dictionary containing the response from the API """ + # Tests the connection before creating it + response = S3Connection.test_connection(client, params) + + logging.info(f"Connection test response: {response}") + unique_id = str(uuid.uuid4()) url = f"{constants.BASE_URL}/connectors/connections/create?client_id={client.client_id}&uuid={unique_id}" @@ -95,7 +125,6 @@ def create_connection( request_payload = { "credentials": (None, creds_payload), "connector": (None, "s3"), - "path": (None, params.s3_path), "connection_type": (None, params.connection_type.value), "name": (None, params.name), "description": (None, params.description), @@ -104,7 +133,8 @@ def create_connection( response_data = client_utils.request( "POST", url, headers=headers, files=request_payload, request_id=unique_id ) - return response_data.get("response", {}) + logging.info(f"Connection creation response: {response_data}") + return LabellerrConnection(client=client, connection_id=response_data.get("response", {}).get("connection_id")) LabellerrConnectionMeta._register("s3", S3Connection) diff --git a/labellerr/core/constants.py b/labellerr/core/constants.py index 91e1bf9..6eb00b4 100644 --- a/labellerr/core/constants.py +++ b/labellerr/core/constants.py @@ -1,4 +1,4 @@ -BASE_URL = "https://api.labellerr.com" +BASE_URL = "https://api-gateway-722091373895.us-central1.run.app" ALLOWED_ORIGINS = "https://pro.labellerr.com" @@ -42,5 +42,4 @@ "dot", "audio", ] -CONNECTION_TYPES = ["s3", "gcs", "local"] cdn_server_address = "cdn-951134552678.us-central1.run.app:443" diff --git a/labellerr/core/schemas/__init__.py b/labellerr/core/schemas/__init__.py index dee02d0..acec40f 100644 --- a/labellerr/core/schemas/__init__.py +++ b/labellerr/core/schemas/__init__.py @@ -18,7 +18,7 @@ from labellerr.core.schemas.base import DirPathStr, FilePathStr, NonEmptyStr # Connection schemas -from labellerr.core.schemas.connections import ( +from labellerr.core.schemas.connectors import ( AWSConnectionParams, AWSConnectorConfig, DatasetDataType, @@ -26,6 +26,8 @@ GCPConnectorConfig, GCSConnectionParams, AWSConnectionTestParams, + ConnectionType, + ConnectorType, ) # Dataset schemas @@ -86,6 +88,8 @@ "DeleteConnectionParams", "AWSConnectorConfig", "GCPConnectorConfig", + "ConnectorType", + "ConnectionType", "DatasetDataType", # Dataset schemas "UploadFilesParams", diff --git a/labellerr/core/schemas/connections.py b/labellerr/core/schemas/connectors.py similarity index 95% rename from labellerr/core/schemas/connections.py rename to labellerr/core/schemas/connectors.py index 741b460..bd4c8dc 100644 --- a/labellerr/core/schemas/connections.py +++ b/labellerr/core/schemas/connectors.py @@ -25,6 +25,12 @@ class ConnectionType(str, Enum): _IMPORT = "import" _EXPORT = "export" +class ConnectorType(str, Enum): + """Enum for connector types.""" + + _S3 = "s3" + _GCS = "gcs" + _LOCAL = "local" class AWSConnectionTestParams(BaseModel): """Parameters for testing an AWS S3 connection.""" From 99aef213b790d6a00e5d7ef65687ce9ead289ec0 Mon Sep 17 00:00:00 2001 From: Ximi Hoque Date: Thu, 6 Nov 2025 00:36:54 +0530 Subject: [PATCH 02/13] Refactored GCS connection, now working fully --- driver.py | 1 - drivers/connectors.py | 39 +++++- labellerr/core/connectors/__init__.py | 29 ++-- labellerr/core/connectors/connections.py | 62 ++++++--- labellerr/core/connectors/gcs_connection.py | 124 ++++++++++++++---- labellerr/core/connectors/s3_connection.py | 43 ++---- labellerr/core/schemas/__init__.py | 8 +- labellerr/core/schemas/connectors.py | 81 ++++-------- .../integration/test_labellerr_integration.py | 2 +- 9 files changed, 230 insertions(+), 159 deletions(-) diff --git a/driver.py b/driver.py index 45efd98..bc71642 100644 --- a/driver.py +++ b/driver.py @@ -26,7 +26,6 @@ from labellerr.core.autolabel import LabellerrAutoLabel - # if os.getenv("CREATE_DATASET", "").lower() == "true": # from labellerr import schemas # from labellerr.core.datasets import create_dataset diff --git a/drivers/connectors.py b/drivers/connectors.py index f9f114c..f5d51d5 100644 --- a/drivers/connectors.py +++ b/drivers/connectors.py @@ -2,8 +2,22 @@ import os from dotenv import load_dotenv from labellerr.client import LabellerrClient -from labellerr.core.connectors import LabellerrConnection, list_connections, delete_connection, create_connection, AWSConnectionParams -from labellerr.core.schemas import ConnectionType, ConnectorType + +# from labellerr.core.connectors import ( +# LabellerrConnection, +# list_connections, +# delete_connection, +# create_connection, +# AWSConnectionParams, +# LabellerrGCSConnection, +# ) +# from labellerr.core.schemas import ( +# ConnectionType, +# ConnectorType, +# DatasetDataType, +# GCSConnectionTestParams, +# GCSConnectionParams, +# ) # Set logging level to DEBUG logging.basicConfig(level=logging.DEBUG) @@ -26,10 +40,27 @@ client_id=CLIENT_ID, ) +# response = LabellerrGCSConnection.test_connection(client=client, params=GCSConnectionTestParams( +# svc_account_json='labellerr-dev.json', +# path="gs://aws-labellerr-public-datasets/coco2017", +# connection_type=ConnectionType._IMPORT, +# data_type=DatasetDataType.video, +# )) +# print(response) + +# response = create_connection(client=client, connector_type=ConnectorType._GCS, params=GCSConnectionParams( +# svc_account_json='labellerr-dev.json', +# path="gs://aws-labellerr-public-datasets/coco2017", +# connection_type=ConnectionType._IMPORT, +# data_type=DatasetDataType.image, name="GCS Import Test", description="GCS Import Test")) +# print(response) +# connection = LabellerrConnection(client=client, connection_id='8c3dc4b4-e701-4d22-add3-28abc33e13ef') +# response = connection.test(path="gs://aws-labellerr-public-datasets/coco2017", connection_type=ConnectionType._IMPORT, data_type=DatasetDataType.image) +# print(response) # connection = create_connection(client=client, connector_type=ConnectorType._S3, params=AWSConnectionParams( # aws_access_key=os.getenv("AWS_KEY"), # aws_secrets_key=os.getenv("AWS_SECRET"), -# s3_path="s3://amazon-s3-sync-test/labellerr-processed/videos/", # this path is not part of the connection but needed to test the connection on the desired path. +# s3_path="s3://amazon-s3-sync-test/labellerr-processed/videos/", # this path is not part of the connection but needed to test the connection on the desired path. # # This can be dynamically changed while using the connection for creating datasets. # connection_type=ConnectionType._IMPORT, # name="Amazon S3 Import Test", @@ -39,4 +70,4 @@ # connection = LabellerrConnection(client=client, connection_id='2a30c044-57f7-42c7-8290-bab3bbac0ebc') # print('connection type', connection.connection_type) -# print (connection.test(s3_path="s3://amazon-s3-sync-test/labellerr-processed/videos/datasets/", connection_type=ConnectionType._IMPORT)) \ No newline at end of file +# print (connection.test(s3_path="s3://amazon-s3-sync-test/labellerr-processed/videos/datasets/", connection_type=ConnectionType._IMPORT)) diff --git a/labellerr/core/connectors/__init__.py b/labellerr/core/connectors/__init__.py index 688bb91..bf1f298 100644 --- a/labellerr/core/connectors/__init__.py +++ b/labellerr/core/connectors/__init__.py @@ -3,6 +3,7 @@ from .. import constants from ...schemas import AWSConnectionParams, ConnectionType, ConnectorType +from ...schemas import GCSConnectionParams from .connections import LabellerrConnection from .gcs_connection import GCSConnection as LabellerrGCSConnection from .s3_connection import S3Connection as LabellerrS3Connection @@ -16,7 +17,7 @@ def create_connection( client: "LabellerrClient", connector_type: ConnectorType, - params: AWSConnectionParams, + params: AWSConnectionParams | GCSConnectionParams, ) -> LabellerrS3Connection | LabellerrGCSConnection: if connector_type == ConnectorType._S3: return LabellerrS3Connection.create_connection(client, params) @@ -27,10 +28,10 @@ def create_connection( def list_connections( - client: "LabellerrClient", - connector: ConnectorType, - connection_type: ConnectionType = None, - ) -> list[LabellerrGCSConnection | LabellerrS3Connection]: + client: "LabellerrClient", + connector: ConnectorType, + connection_type: ConnectionType = None, +) -> list[LabellerrGCSConnection | LabellerrS3Connection]: """ Lists connections for a client :param client: LabellerrClient instance @@ -52,13 +53,13 @@ def list_connections( extra_headers = {"email_id": client.api_key} response = client.make_request( - "GET", - url, - extra_headers=extra_headers, - request_id=unique_id, - params=params + "GET", url, extra_headers=extra_headers, request_id=unique_id, params=params ) - return [LabellerrConnection(client, connection["connection_id"]) for connection in response.get("response", [])] + return [ + LabellerrConnection(client, connection["connection_id"]) + for connection in response.get("response", []) + ] + def delete_connection(client: "LabellerrClient", connection_id: str): """ @@ -75,6 +76,10 @@ def delete_connection(client: "LabellerrClient", connection_id: str): extra_headers = {"email_id": client.api_key} response = client.make_request( - "POST", url, extra_headers=extra_headers, request_id=request_id, json={"connection_id": connection_id} + "POST", + url, + extra_headers=extra_headers, + request_id=request_id, + json={"connection_id": connection_id}, ) return response.get("response", None) diff --git a/labellerr/core/connectors/connections.py b/labellerr/core/connectors/connections.py index 718ea22..c691813 100644 --- a/labellerr/core/connectors/connections.py +++ b/labellerr/core/connectors/connections.py @@ -1,11 +1,10 @@ """This module will contain all CRUD for connections. Example, create, list connections, get connection, delete connection, update connection, etc.""" import uuid -from abc import ABCMeta, abstractmethod +from abc import ABCMeta from typing import TYPE_CHECKING, Dict -import logging from .. import client_utils, constants -from ..schemas import ConnectionType, ConnectorType +from ..schemas import ConnectionType, DatasetDataType from ..exceptions import InvalidConnectionError, InvalidDatasetIDError if TYPE_CHECKING: @@ -28,20 +27,13 @@ def get_connection(client: "LabellerrClient", connection_id: str): assert connection_id, "Connection ID is can't be empty" unique_id = str(uuid.uuid4()) url = f"{constants.BASE_URL}/connectors/connections/{connection_id}/details" - - params = { - "client_id": client.client_id, - "uuid": unique_id - } - + + params = {"client_id": client.client_id, "uuid": unique_id} + extra_headers = {"content-type": "application/json"} - + response = client.make_request( - "GET", - url, - extra_headers=extra_headers, - request_id=unique_id, - params=params + "GET", url, extra_headers=extra_headers, request_id=unique_id, params=params ) return response.get("response", None) @@ -89,11 +81,11 @@ def connection_id(self): @property def connection_type(self): return self.__connection_data.get("connection_type") - + @property def connector(self): return self.__connection_data.get("connector") - + @property def created_at(self): return self.__connection_data.get("created_at") @@ -102,7 +94,35 @@ def created_at(self): def created_by(self): return self.__connection_data.get("created_by") - @abstractmethod - def test(self): - """Each connection type must implement its own connection testing logic""" - pass \ No newline at end of file + def test( + self, path: str, connection_type: ConnectionType, data_type: DatasetDataType + ): + request_id = str(uuid.uuid4()) + test_connection_url = ( + f"{constants.BASE_URL}/connectors/connections/test" + f"?client_id={self.client.client_id}&uuid={request_id}" + ) + + headers = client_utils.build_headers( + api_key=self.client.api_key, + api_secret=self.client.api_secret, + client_id=self.client.client_id, + extra_headers={"email_id": self.client.api_key}, + ) + + # Test endpoint also expects multipart/form-data format + test_request = { + "connector": (None, self.connector), + "path": (None, path), + "connection_type": (None, connection_type.value), + "connection_id": (None, self.connection_id), + "data_type": (None, data_type.value), + } + response = client_utils.request( + "POST", + test_connection_url, + headers=headers, + files=test_request, + request_id=request_id, + ) + return response.get("response", {}) diff --git a/labellerr/core/connectors/gcs_connection.py b/labellerr/core/connectors/gcs_connection.py index 6ff811f..da114b2 100644 --- a/labellerr/core/connectors/gcs_connection.py +++ b/labellerr/core/connectors/gcs_connection.py @@ -1,21 +1,72 @@ import uuid -from typing import TYPE_CHECKING - +import logging +import os from .. import client_utils, constants +from ..client import LabellerrClient from .connections import LabellerrConnection, LabellerrConnectionMeta - -if TYPE_CHECKING: - from labellerr import LabellerrClient +from ..schemas import GCSConnectionTestParams, GCSConnectionParams +from ..exceptions import LabellerrError class GCSConnection(LabellerrConnection): - def test_connection(self): - print("Testing GCS connection!") - return True + @staticmethod + def test_connection( + client: "LabellerrClient", params: GCSConnectionTestParams + ) -> dict: + request_id = str(uuid.uuid4()) + url = ( + f"{constants.BASE_URL}/connectors/connections/test" + f"?client_id={client.client_id}&uuid={request_id}" + ) + + # Prepare multipart form data + files = [] + if params.svc_account_json and os.path.exists(params.svc_account_json): + # If file path is provided, read and upload as file + with open(params.svc_account_json, "rb") as f: + files = [ + ( + "attachment_files", + ( + params.svc_account_json.split("/")[-1], + f.read(), + "application/json", + ), + ) + ] + else: + raise LabellerrError("Service account JSON file is required") + + # Prepare form data payload + form_data = { + "connector": "gcs", + "path": params.path, + "connection_type": params.connection_type.value, + "data_type": params.data_type.value, + } + + # Build headers without content-type for multipart form data + headers = client_utils.build_headers( + api_key=client.api_key, + api_secret=client.api_secret, + client_id=client.client_id, + ) + + response = client_utils.request( + "POST", + url, + headers=headers, + data=form_data, + files=files, + request_id=request_id, + ) + return response.get("response", {}) @staticmethod - def create_connection(client: "LabellerrClient", gcp_config: dict) -> str: + def create_connection( + client: "LabellerrClient", params: GCSConnectionParams + ) -> str: """ Sets up GCP connector for dataset creation (quick connection). @@ -23,37 +74,54 @@ def create_connection(client: "LabellerrClient", gcp_config: dict) -> str: :param gcp_config: GCP configuration containing bucket_name, folder_path, service_account_key :return: Connection ID for GCP connector """ - import json - - from ... import LabellerrError - - required_fields = ["bucket_name"] - for field in required_fields: - if field not in gcp_config: - raise LabellerrError(f"Required field '{field}' missing in gcp_config") + response = GCSConnection.test_connection(client, params) + logging.info(f"GCS connection test response: {response}") unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/connectors/connect/gcp?client_id={client.client_id}&uuid={unique_id}" + url = f"{constants.BASE_URL}/connectors/connections/create?client_id={client.client_id}&uuid={unique_id}" + # Build headers without content-type for multipart form data headers = client_utils.build_headers( api_key=client.api_key, api_secret=client.api_secret, client_id=client.client_id, - extra_headers={"content-type": "application/json"}, ) - payload = json.dumps( - { - "bucket_name": gcp_config["bucket_name"], - "folder_path": gcp_config.get("folder_path", ""), - "service_account_key": gcp_config.get("service_account_key"), - } - ) + # Prepare multipart form data + with open(params.svc_account_json, "rb") as f: + files = [ + ( + "attachment_files", + ( + params.svc_account_json.split("/")[-1], + f.read(), + "application/json", + ), + ) + ] + + # Prepare form data (not JSON) + form_data = { + "connector": "gcs", + "connection_type": params.connection_type.value, + "name": params.name, + "description": params.description, + "credentials": "svc_account_json", + "client_id": client.client_id, + } response_data = client_utils.request( - "POST", url, headers=headers, data=payload, request_id=unique_id + "POST", + url, + headers=headers, + data=form_data, + files=files, + request_id=unique_id, + ) + return LabellerrConnection( + client=client, + connection_id=response_data.get("response", {}).get("connection_id"), ) - return response_data["response"]["connection_id"] LabellerrConnectionMeta._register("gcs", GCSConnection) diff --git a/labellerr/core/connectors/s3_connection.py b/labellerr/core/connectors/s3_connection.py index d65095b..8aa15b4 100644 --- a/labellerr/core/connectors/s3_connection.py +++ b/labellerr/core/connectors/s3_connection.py @@ -1,10 +1,12 @@ import json import uuid from ..client import LabellerrClient -from ..schemas import AWSConnectionParams, AWSConnectionTestParams, ConnectionType +from ..schemas import AWSConnectionParams, AWSConnectionTestParams from .. import client_utils, constants from .connections import LabellerrConnection, LabellerrConnectionMeta import logging + + class S3Connection(LabellerrConnection): @staticmethod @@ -41,8 +43,9 @@ def test_connection( test_request = { "credentials": (None, aws_credentials_json), "connector": (None, "s3"), - "path": (None, params.s3_path), + "path": (None, params.path), "connection_type": (None, params.connection_type.value), + "data_type": (None, params.data_type.value), } # Remove content-type from headers to let requests set it with boundary @@ -59,36 +62,6 @@ def test_connection( ) return response.get("response", {}) - def test(self, s3_path: str, connection_type: ConnectionType): - request_id = str(uuid.uuid4()) - test_connection_url = ( - f"{constants.BASE_URL}/connectors/connections/test" - f"?client_id={self.client.client_id}&uuid={request_id}" - ) - - headers = client_utils.build_headers( - api_key=self.client.api_key, - api_secret=self.client.api_secret, - client_id=self.client.client_id, - extra_headers={"email_id": self.client.api_key}, - ) - - # Test endpoint also expects multipart/form-data format - test_request = { - "connector": (None, "s3"), - "path": (None, s3_path), - "connection_type": (None, connection_type.value), - "connection_id": (None, self.connection_id), - } - response = client_utils.request( - "POST", - test_connection_url, - headers=headers, - files=test_request, - request_id=request_id, - ) - return response.get("response", {}) - @staticmethod def create_connection( client: "LabellerrClient", params: AWSConnectionParams @@ -133,8 +106,10 @@ def create_connection( response_data = client_utils.request( "POST", url, headers=headers, files=request_payload, request_id=unique_id ) - logging.info(f"Connection creation response: {response_data}") - return LabellerrConnection(client=client, connection_id=response_data.get("response", {}).get("connection_id")) + return LabellerrConnection( + client=client, + connection_id=response_data.get("response", {}).get("connection_id"), + ) LabellerrConnectionMeta._register("s3", S3Connection) diff --git a/labellerr/core/schemas/__init__.py b/labellerr/core/schemas/__init__.py index acec40f..9fa4967 100644 --- a/labellerr/core/schemas/__init__.py +++ b/labellerr/core/schemas/__init__.py @@ -20,14 +20,13 @@ # Connection schemas from labellerr.core.schemas.connectors import ( AWSConnectionParams, - AWSConnectorConfig, DatasetDataType, DeleteConnectionParams, - GCPConnectorConfig, GCSConnectionParams, AWSConnectionTestParams, ConnectionType, ConnectorType, + GCSConnectionTestParams, ) # Dataset schemas @@ -85,11 +84,10 @@ "AWSConnectionParams", "AWSConnectionTestParams", "GCSConnectionParams", + "GCSConnectionTestParams", "DeleteConnectionParams", - "AWSConnectorConfig", - "GCPConnectorConfig", "ConnectorType", - "ConnectionType", + "ConnectionType", "DatasetDataType", # Dataset schemas "UploadFilesParams", diff --git a/labellerr/core/schemas/connectors.py b/labellerr/core/schemas/connectors.py index bd4c8dc..eb1485e 100644 --- a/labellerr/core/schemas/connectors.py +++ b/labellerr/core/schemas/connectors.py @@ -4,7 +4,7 @@ import os from enum import Enum -from typing import Literal, Optional +from typing import Optional from pydantic import BaseModel, Field, field_validator @@ -25,6 +25,7 @@ class ConnectionType(str, Enum): _IMPORT = "import" _EXPORT = "export" + class ConnectorType(str, Enum): """Enum for connector types.""" @@ -32,40 +33,45 @@ class ConnectorType(str, Enum): _GCS = "gcs" _LOCAL = "local" -class AWSConnectionTestParams(BaseModel): - """Parameters for testing an AWS S3 connection.""" - aws_access_key: str = Field(min_length=1) - aws_secrets_key: str = Field(min_length=1) - s3_path: str = Field(min_length=1) +class GCSConnectionTestParams(BaseModel): + """Parameters for testing a GCS connection.""" + + svc_account_json: Optional[str] = Field(default=None, min_length=1) + path: str = Field(min_length=1) connection_type: ConnectionType = ConnectionType._IMPORT + data_type: DatasetDataType + @field_validator("svc_account_json") + @classmethod + def validate_svc_account_json(cls, v): + if v and not os.path.exists(v): + raise ValueError(f"GCS credential file not found: {v}") + return v -class AWSConnectionParams(AWSConnectionTestParams): - """Parameters for creating an AWS S3 connection.""" + +class GCSConnectionParams(GCSConnectionTestParams): + """Parameters for creating a GCS connection.""" name: str = Field(min_length=1) description: str -class GCSConnectionParams(BaseModel): - """Parameters for creating a GCS connection.""" +class AWSConnectionTestParams(BaseModel): + """Parameters for testing an AWS S3 connection.""" - client_id: str = Field(min_length=1) - gcs_cred_file: str - gcs_path: str = Field(min_length=1) + aws_access_key: str = Field(min_length=1) + aws_secrets_key: str = Field(min_length=1) + path: str = Field(min_length=1) + connection_type: ConnectionType = ConnectionType._IMPORT data_type: DatasetDataType + + +class AWSConnectionParams(AWSConnectionTestParams): + """Parameters for creating an AWS S3 connection.""" + name: str = Field(min_length=1) description: str - connection_type: str = "import" - credentials: str = "svc_account_json" - - @field_validator("gcs_cred_file") - @classmethod - def validate_gcs_cred_file(cls, v): - if not os.path.exists(v): - raise ValueError(f"GCS credential file not found: {v}") - return v class DeleteConnectionParams(BaseModel): @@ -73,34 +79,3 @@ class DeleteConnectionParams(BaseModel): client_id: str = Field(min_length=1) connection_id: str = Field(min_length=1) - - -class AWSConnectorConfig(BaseModel): - """Configuration for AWS S3 connector.""" - - aws_access_key: str = Field(min_length=1) - aws_secrets_key: str = Field(min_length=1) - s3_path: str = Field(min_length=1) - data_type: Literal["image", "video", "audio", "document", "text"] - name: Optional[str] = None - description: str = "Auto-created AWS connector" - connection_type: str = "import" - - -class GCPConnectorConfig(BaseModel): - """Configuration for GCP connector.""" - - gcs_cred_file: str = Field(min_length=1) - gcs_path: str = Field(min_length=1) - data_type: Literal["image", "video", "audio", "document", "text"] - name: Optional[str] = None - description: str = "Auto-created GCS connector" - connection_type: str = "import" - credentials: str = "svc_account_json" - - @field_validator("gcs_cred_file") - @classmethod - def validate_gcs_cred_file(cls, v): - if not os.path.exists(v): - raise ValueError(f"GCS credential file not found: {v}") - return v diff --git a/tests/integration/test_labellerr_integration.py b/tests/integration/test_labellerr_integration.py index f19ff48..9b12513 100644 --- a/tests/integration/test_labellerr_integration.py +++ b/tests/integration/test_labellerr_integration.py @@ -400,7 +400,7 @@ def test_aws_connection_lifecycle(self, integration_client, test_credentials): client_id=test_credentials["client_id"], aws_access_key=aws_secret.get("access_key"), aws_secrets_key=aws_secret.get("secret_key"), - s3_path=aws_secret.get("s3_path"), + path=aws_secret.get("s3_path"), data_type=DatasetDataType.image, name=connection_name, description="Test AWS connection", From 1d610fa038777463667102006a693506bb3da0fc Mon Sep 17 00:00:00 2001 From: Ximi Hoque Date: Thu, 6 Nov 2025 18:30:42 +0530 Subject: [PATCH 03/13] Create dataset refactored, list datasets --- drivers/datasets.py | 62 +++++ labellerr/core/datasets/__init__.py | 286 ++++++++++++++---------- labellerr/core/datasets/base.py | 192 +++++++--------- labellerr/core/datasets/utils.py | 22 +- labellerr/core/schemas/base.py | 11 + labellerr/core/schemas/connectors.py | 12 +- labellerr/core/schemas/datasets.py | 6 +- tests/integration/test_sync_datasets.py | 10 +- 8 files changed, 339 insertions(+), 262 deletions(-) create mode 100644 drivers/datasets.py diff --git a/drivers/datasets.py b/drivers/datasets.py new file mode 100644 index 0000000..d89fcf4 --- /dev/null +++ b/drivers/datasets.py @@ -0,0 +1,62 @@ +from labellerr.core.datasets import ( + create_dataset_from_local, + create_dataset_from_connection, +) +import logging +import os +from dotenv import load_dotenv +from labellerr.client import LabellerrClient +from labellerr.core.schemas import DatasetConfig +from labellerr.core.datasets import LabellerrDataset + +# from labellerr.core.connectors import ( +# LabellerrConnection, +# list_connections, +# delete_connection, +# create_connection, +# AWSConnectionParams, +# LabellerrGCSConnection, +# ) +# from labellerr.core.schemas import ( +# ConnectionType, +# ConnectorType, +# DatasetDataType, +# GCSConnectionTestParams, +# GCSConnectionParams, +# ) + +# Set logging level to DEBUG +logging.basicConfig(level=logging.DEBUG) + +load_dotenv() + +API_KEY = os.getenv("API_KEY") +API_SECRET = os.getenv("API_SECRET") +CLIENT_ID = os.getenv("CLIENT_ID") + +if not all([API_KEY, API_SECRET, CLIENT_ID]): + raise ValueError( + "API_KEY, API_SECRET, and CLIENT_ID must be set in environment variables" + ) + +# Initialize client +client = LabellerrClient( + api_key=API_KEY, + api_secret=API_SECRET, + client_id=CLIENT_ID, +) +dataset_config = DatasetConfig( + dataset_name="test_dataset_from_local SDK", + dataset_description="test dataset description", + data_type="image", +) +# dataset = create_dataset_from_local( +# client=client, +# dataset_config=dataset_config, +# folder_to_upload='images_single', +# ) +dataset = LabellerrDataset( + client=client, dataset_id="455e3d45-55f9-436d-98c2-07a514b7894e" +) +print(dataset.dataset_data) +print(dataset.files_count) diff --git a/labellerr/core/datasets/__init__.py b/labellerr/core/datasets/__init__.py index 28c782c..c7d2707 100644 --- a/labellerr/core/datasets/__init__.py +++ b/labellerr/core/datasets/__init__.py @@ -3,9 +3,7 @@ import uuid from typing import TYPE_CHECKING -from ... import schemas as root_schemas from .. import constants, schemas -from ..connectors import create_connection from ..exceptions import LabellerrError from .audio_dataset import AudioDataSet as LabellerrAudioDataset from .base import LabellerrDataset @@ -26,14 +24,51 @@ ] -def create_dataset( +def create_dataset_from_connection( + client: "LabellerrClient", + dataset_config: schemas.DatasetConfig, + connection_id: str, + path: str, +) -> LabellerrDataset: + """ + Creates a dataset via a connection. + + :param client: The client to use for the request. + :param dataset_config: The configuration for the dataset. + :param connection_id: The ID of the connection to use for the dataset. + :param path: The path to the data source. + :return: The LabellerrDataset instance. + """ + unique_id = str(uuid.uuid4()) + url = f"{constants.BASE_URL}/datasets/create?client_id={client.client_id}&uuid={unique_id}" + + payload = json.dumps( + { + "dataset_name": dataset_config.dataset_name, + "dataset_description": dataset_config.dataset_description, + "data_type": dataset_config.data_type, + "connection_id": connection_id, + "path": path, + "client_id": client.client_id, + "es_multimodal_index": dataset_config.multimodal_indexing, + } + ) + response_data = client.make_request( + "POST", + url, + extra_headers={"content-type": "application/json"}, + request_id=unique_id, + data=payload, + ) + dataset_id = response_data["response"]["dataset_id"] + return LabellerrDataset(client=client, dataset_id=dataset_id) # type: ignore[abstract] + + +def create_dataset_from_local( client: "LabellerrClient", dataset_config: schemas.DatasetConfig, files_to_upload=None, folder_to_upload=None, - path=None, - connection_id=None, - connector_config=None, ): """ Creates a dataset with support for multiple data types and connectors. @@ -44,129 +79,138 @@ def create_dataset( Can also be a DatasetConfig Pydantic model instance. :param files_to_upload: List of file paths to upload (for local connector) :param folder_to_upload: Path to folder to upload (for local connector) - :param connection_id: Pre-existing connection ID to use for the dataset. - Either connection_id or connector_config can be provided, but not both. - :param connector_config: Configuration for cloud connectors (GCP/AWS) - Can be a dict or AWSConnectorConfig/GCPConnectorConfig model instance. - Either connection_id or connector_config can be provided, but not both. - :return: A dictionary containing the response status and the ID of the created dataset. - :raises LabellerrError: If both connection_id and connector_config are provided. + :return: The LabellerrDataset instance. + """ + if files_to_upload is not None: + connection_id = upload_files( + client, + client_id=client.client_id, + files_list=files_to_upload, + ) + elif folder_to_upload is not None: + result = upload_folder_files_to_dataset( + client, + { + "client_id": client.client_id, + "folder_path": folder_to_upload, + "data_type": dataset_config.data_type, + }, + ) + connection_id = result.pop("connection_id") + logging.info(f"Folder uploaded successfully. {result}") + else: + raise LabellerrError("No files or folder to upload provided") + + return create_dataset_from_connection( + client, + dataset_config, + connection_id, + path="local", + ) + + +def delete_dataset(client: "LabellerrClient", dataset_id: str): + """ + Deletes a dataset from the system. + + :param dataset_id: The ID of the dataset to delete + :return: Dictionary containing deletion status + :raises LabellerrError: If the deletion fails + """ + unique_id = str(uuid.uuid4()) + url = f"{constants.BASE_URL}/datasets/{dataset_id}/delete?client_id={client.client_id}&uuid={unique_id}" + + return client.make_request( + "DELETE", + url, + extra_headers={"content-type": "application/json"}, + request_id=unique_id, + ) + + +def list_datasets( + client: "LabellerrClient", + datatype: str, + scope: schemas.DataSetScope, + page_size: int = constants.DEFAULT_PAGE_SIZE, + last_dataset_id: str = None, +): """ + Retrieves datasets by parameters with pagination support. + Always returns a generator that yields individual datasets. + + :param client: The client object. + :param datatype: The type of data for the dataset. + :param scope: The permission scope for the dataset. + :param page_size: Number of datasets to return per page (default: 10) + Use -1 to auto-paginate through all pages + Use specific number to fetch only that many datasets from first page + :param last_dataset_id: ID of the last dataset from previous page for pagination + (only used when page_size is a specific number, ignored for -1) + :return: Generator yielding individual datasets + + Examples: + # Auto-paginate through all datasets + for dataset in get_all_datasets(client, "image", DataSetScope.client, page_size=-1): + print(dataset) + + # Get first 20 datasets + datasets = list(get_all_datasets(client, "image", DataSetScope.client, page_size=20)) + + # Manual pagination - first page of 10 + gen = get_all_datasets(client, "image", DataSetScope.client, page_size=10) + first_10 = list(gen) + """ + # Auto-pagination mode: yield datasets across all pages + if page_size == -1: + actual_page_size = constants.DEFAULT_PAGE_SIZE + current_last_dataset_id = None + has_more = True + + while has_more: + unique_id = str(uuid.uuid4()) + url = ( + f"{constants.BASE_URL}/datasets/list?client_id={client.client_id}&data_type={datatype}&permission_level={scope}" + f"&page_size={actual_page_size}&uuid={unique_id}" + ) - try: - # Validate that both connection_id and connector_config are not provided - if connection_id is not None and connector_config is not None: - raise LabellerrError( - "Cannot provide both connection_id and connector_config. " - "Use connection_id for existing connections or connector_config to create a new connection." + if current_last_dataset_id: + url += f"&last_dataset_id={current_last_dataset_id}" + + response = client.make_request( + "GET", + url, + extra_headers={"content-type": "application/json"}, + request_id=unique_id, ) - connector_type = dataset_config.connector_type - # Use provided connection_id or set to None (will be created later if needed) - final_connection_id = connection_id - - # Handle different connector types only if connection_id is not provided - if final_connection_id is None: - if connector_type == "local": - if files_to_upload is not None: - try: - final_connection_id = upload_files( - client, - client_id=client.client_id, - files_list=files_to_upload, - ) - except Exception as e: - raise LabellerrError( - f"Failed to upload files to dataset: {str(e)}" - ) - - elif folder_to_upload is not None: - try: - result = upload_folder_files_to_dataset( - client, - { - "client_id": client.client_id, - "folder_path": folder_to_upload, - "data_type": dataset_config.data_type, - }, - ) - final_connection_id = result["connection_id"] - except Exception as e: - raise LabellerrError( - f"Failed to upload folder files to dataset: {str(e)}" - ) - elif connector_config is None: - # Create empty dataset for local connector - final_connection_id = None - - elif connector_type in ["gcp", "aws"]: - if connector_config is None: - raise LabellerrError( - f"connector_config is required for {connector_type} connector when connection_id is not provided" - ) - if path is None: - raise LabellerrError( - f"path is required for {connector_type} connector" - ) - - # Validate connector_config using Pydantic models - if connector_type == "aws": - if not isinstance( - connector_config, root_schemas.AWSConnectorConfig - ): - validated_connector = root_schemas.AWSConnectorConfig( - **connector_config - ) - else: - validated_connector = connector_config - else: # gcp - if not isinstance( - connector_config, root_schemas.GCPConnectorConfig - ): - validated_connector = root_schemas.GCPConnectorConfig( - **connector_config - ) - else: - validated_connector = connector_config - - try: - final_connection_id = create_connection( - client, - connector_type, - client.client_id, - validated_connector.model_dump(), - ) - except Exception as e: - raise LabellerrError( - f"Failed to setup {connector_type} connector: {str(e)}" - ) - else: - raise LabellerrError(f"Unsupported connector type: {connector_type}") + datasets = response.get("response", {}).get("datasets", []) + for dataset in datasets: + yield dataset - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/datasets/create?client_id={client.client_id}&uuid={unique_id}" + # Check if there are more pages + has_more = response.get("response", {}).get("has_more", False) + current_last_dataset_id = response.get("response", {}).get( + "last_dataset_id" + ) - payload = json.dumps( - { - "dataset_name": dataset_config.dataset_name, - "dataset_description": dataset_config.dataset_description, - "data_type": dataset_config.data_type, - "connection_id": final_connection_id, - "path": path, - "client_id": client.client_id, - "connector_type": connector_type, - } + else: + unique_id = str(uuid.uuid4()) + url = ( + f"{constants.BASE_URL}/datasets/list?client_id={client.client_id}&data_type={datatype}&permission_level={scope}" + f"&page_size={page_size}&uuid={unique_id}" ) - response_data = client.make_request( - "POST", + + # Add last_dataset_id for pagination if provided + if last_dataset_id: + url += f"&last_dataset_id={last_dataset_id}" + + response = client.make_request( + "GET", url, extra_headers={"content-type": "application/json"}, request_id=unique_id, - data=payload, ) - dataset_id = response_data["response"]["dataset_id"] - return LabellerrDataset(client=client, dataset_id=dataset_id) # type: ignore[abstract] - - except LabellerrError as e: - logging.error(f"Failed to create dataset: {e}") - raise + datasets = response.get("response", {}).get("datasets", []) + for dataset in datasets: + yield dataset diff --git a/labellerr/core/datasets/base.py b/labellerr/core/datasets/base.py index 177adc3..d980637 100644 --- a/labellerr/core/datasets/base.py +++ b/labellerr/core/datasets/base.py @@ -1,16 +1,15 @@ """This module will contain all CRUD for datasets. Example, create, list datasets, get dataset, delete dataset, update dataset, etc.""" import json +import logging import uuid from abc import ABCMeta, abstractmethod -from typing import TYPE_CHECKING, Dict +from typing import Dict, Optional, Any from ...schemas import DataSetScope from .. import constants from ..exceptions import InvalidDatasetError - -if TYPE_CHECKING: - from ..client import LabellerrClient +from ..client import LabellerrClient class LabellerrDatasetMeta(ABCMeta): @@ -69,136 +68,117 @@ class LabellerrDataset(metaclass=LabellerrDatasetMeta): def __init__(self, client: "LabellerrClient", dataset_id: str, **kwargs): self.client = client self.dataset_id = dataset_id - self.dataset_data = kwargs["dataset_data"] + self.__dataset_data = kwargs["dataset_data"] @property - def files_count(self): - return self.dataset_data.get("files_count", 0) + def name(self): + return self.__dataset_data.get("name") @property - def status_code(self): - return self.dataset_data.get("status_code", 501) # if not found, return 501 + def description(self): + return self.__dataset_data.get("description") @property - def data_type(self): - return self.dataset_data.get("data_type") + def created_at(self): + return self.__dataset_data.get("created_at") - @abstractmethod - def fetch_files(self): - """Each file type must implement its own download logic""" - pass + @property + def created_by(self): + return self.__dataset_data.get("created_by") - @staticmethod - def get_all_datasets( - client: "LabellerrClient", - datatype: str, - scope: DataSetScope, - page_size: int = None, - last_dataset_id: str = None, - ): - """ - Retrieves datasets by parameters with pagination support. - Always returns a generator that yields individual datasets. - - :param client: The client object. - :param datatype: The type of data for the dataset. - :param scope: The permission scope for the dataset. - :param page_size: Number of datasets to return per page (default: 10) - Use -1 to auto-paginate through all pages - Use specific number to fetch only that many datasets from first page - :param last_dataset_id: ID of the last dataset from previous page for pagination - (only used when page_size is a specific number, ignored for -1) - :return: Generator yielding individual datasets + @property + def files_count(self): + return self.__dataset_data.get("files_count", 0) - Examples: - # Auto-paginate through all datasets - for dataset in get_all_datasets(client, "image", DataSetScope.client, page_size=-1): - print(dataset) + @property + def status_code(self): + return self.__dataset_data.get("status_code", 501) # if not found, return 501 - # Get first 20 datasets - datasets = list(get_all_datasets(client, "image", DataSetScope.client, page_size=20)) + @property + def data_type(self): + return self.__dataset_data.get("data_type") - # Manual pagination - first page of 10 - gen = get_all_datasets(client, "image", DataSetScope.client, page_size=10) - first_10 = list(gen) + def status( + self, + interval: float = 2.0, + timeout: Optional[float] = None, + max_retries: Optional[int] = None, + ) -> Dict[str, Any]: """ - # Set default page size if not specified - if page_size is None: - page_size = constants.DEFAULT_PAGE_SIZE - - # Auto-pagination mode: yield datasets across all pages - if page_size == -1: - actual_page_size = constants.DEFAULT_PAGE_SIZE - current_last_dataset_id = None - has_more = True - - while has_more: - unique_id = str(uuid.uuid4()) - url = ( - f"{constants.BASE_URL}/datasets/list?client_id={client.client_id}&data_type={datatype}&permission_level={scope}" - f"&page_size={actual_page_size}&uuid={unique_id}" - ) + Poll dataset status until completion or timeout. - if current_last_dataset_id: - url += f"&last_dataset_id={current_last_dataset_id}" + Args: + interval: Time in seconds between status checks (default: 2.0) + timeout: Maximum time in seconds to poll before giving up + max_retries: Maximum number of retries before giving up - response = client.make_request( - "GET", - url, - extra_headers={"content-type": "application/json"}, - request_id=unique_id, - ) + Returns: + Final dataset data with status information - datasets = response.get("response", {}).get("datasets", []) - for dataset in datasets: - yield dataset + Examples: + # Poll until dataset processing is complete + final_status = dataset.status() - # Check if there are more pages - has_more = response.get("response", {}).get("has_more", False) - current_last_dataset_id = response.get("response", {}).get( - "last_dataset_id" - ) + # Poll with custom timeout + final_status = dataset.status(timeout=300) + + # Poll with custom interval and max retries + final_status = dataset.status(interval=5.0, max_retries=20) + """ + from ..utils import poll - else: + def get_dataset_status(): unique_id = str(uuid.uuid4()) url = ( - f"{constants.BASE_URL}/datasets/list?client_id={client.client_id}&data_type={datatype}&permission_level={scope}" - f"&page_size={page_size}&uuid={unique_id}" + f"{constants.BASE_URL}/datasets/{self.dataset_id}?client_id={self.client.client_id}" + f"&uuid={unique_id}" ) - # Add last_dataset_id for pagination if provided - if last_dataset_id: - url += f"&last_dataset_id={last_dataset_id}" - - response = client.make_request( + response = self.client.make_request( "GET", url, extra_headers={"content-type": "application/json"}, request_id=unique_id, ) - datasets = response.get("response", {}).get("datasets", []) - for dataset in datasets: - yield dataset - - def delete_dataset(self, dataset_id): - """ - Deletes a dataset from the system. - - :param dataset_id: The ID of the dataset to delete - :return: Dictionary containing deletion status - :raises LabellerrError: If the deletion fails - """ - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/datasets/{dataset_id}/delete?client_id={self.client.client_id}&uuid={unique_id}" - - return self.client.make_request( - "DELETE", - url, - extra_headers={"content-type": "application/json"}, - request_id=unique_id, + dataset_data = response.get("response", {}) + if dataset_data: + self.__dataset_data = dataset_data + return dataset_data + + def is_completed(dataset_data): + status_code = dataset_data.get("status_code", 500) + # Consider dataset complete when status_code is 200 (success) or >= 400 (error/failed) + return status_code == 200 or status_code >= 400 + + def on_success(dataset_data): + status_code = dataset_data.get("status_code", 500) + if status_code == 300: + logging.info( + "Dataset %s processing completed successfully!", self.dataset_id + ) + else: + logging.warning( + "Dataset %s processing finished with status code: %s", + self.dataset_id, + status_code, + ) + return dataset_data + + return poll( + function=get_dataset_status, + condition=is_completed, + interval=interval, + timeout=timeout, + max_retries=max_retries, + on_success=on_success, ) - def sync_datasets( + @abstractmethod + def fetch_files(self): + """Each file type must implement its own download logic""" + pass + + def sync_with_connection( self, project_id, path, diff --git a/labellerr/core/datasets/utils.py b/labellerr/core/datasets/utils.py index a8fe268..2be04cb 100644 --- a/labellerr/core/datasets/utils.py +++ b/labellerr/core/datasets/utils.py @@ -110,7 +110,6 @@ def connect_local_files( return client_utils.request("POST", url, headers=headers, json=body) -@validate_params(client_id=str, files_list=(str, list)) def upload_files( client: "LabellerrClient", client_id: str, files_list: Union[str, List[str]] ): @@ -122,23 +121,12 @@ def upload_files( :return: The connection ID from the API. :raises LabellerrError: If the upload fails. """ - # Validate parameters using Pydantic - params = schemas.UploadFilesParams(client_id=client_id, files_list=files_list) - try: - # Use validated files_list from Pydantic - files_list = params.files_list - - if len(files_list) == 0: - raise LabellerrError("No files to upload") + if len(files_list) == 0: + raise LabellerrError("No files to upload") - response = __process_batch(client, client_id, files_list) - connection_id = response["response"]["temporary_connection_id"] - return connection_id - except LabellerrError: - raise - except Exception as e: - logging.error(f"Failed to upload files: {str(e)}") - raise + response = __process_batch(client, client_id, files_list) + connection_id = response["response"]["temporary_connection_id"] + return connection_id def __process_batch( diff --git a/labellerr/core/schemas/base.py b/labellerr/core/schemas/base.py index 39cde3e..789c68a 100644 --- a/labellerr/core/schemas/base.py +++ b/labellerr/core/schemas/base.py @@ -3,6 +3,17 @@ """ import os +from enum import Enum + + +class DatasetDataType(str, Enum): + """Enum for dataset data types.""" + + image = "image" + video = "video" + audio = "audio" + document = "document" + text = "text" class NonEmptyStr(str): diff --git a/labellerr/core/schemas/connectors.py b/labellerr/core/schemas/connectors.py index eb1485e..aa3eaf4 100644 --- a/labellerr/core/schemas/connectors.py +++ b/labellerr/core/schemas/connectors.py @@ -5,20 +5,10 @@ import os from enum import Enum from typing import Optional - +from .base import DatasetDataType from pydantic import BaseModel, Field, field_validator -class DatasetDataType(str, Enum): - """Enum for dataset data types.""" - - image = "image" - video = "video" - audio = "audio" - document = "document" - text = "text" - - class ConnectionType(str, Enum): """Enum for connection types.""" diff --git a/labellerr/core/schemas/datasets.py b/labellerr/core/schemas/datasets.py index 6901a3c..555fbbd 100644 --- a/labellerr/core/schemas/datasets.py +++ b/labellerr/core/schemas/datasets.py @@ -7,6 +7,8 @@ from typing import List, Literal from uuid import UUID +from .connectors import ConnectorType +from .base import DatasetDataType from pydantic import BaseModel, Field, field_validator @@ -107,6 +109,6 @@ class DatasetConfig(BaseModel): """Configuration for creating a dataset.""" dataset_name: str = Field(min_length=1) - data_type: Literal["image", "video", "audio", "document", "text"] + data_type: DatasetDataType dataset_description: str = "" - connector_type: Literal["local", "aws", "gcp"] = "local" + multimodal_indexing: bool = False diff --git a/tests/integration/test_sync_datasets.py b/tests/integration/test_sync_datasets.py index 902b76e..bdd7db2 100644 --- a/tests/integration/test_sync_datasets.py +++ b/tests/integration/test_sync_datasets.py @@ -89,7 +89,7 @@ def test_sync_datasets_aws(self): print(f"Data Type: {self.data_type}") print(f"Email ID: {self.email_id}") - response = datasets.sync_datasets( + response = datasets.sync_with_connection( client_id=self.client_id, project_id=self.project_id, dataset_id=self.aws_dataset_id, @@ -134,7 +134,7 @@ def test_sync_datasets_gcs(self): print(f"Data Type: {self.data_type}") print(f"Email ID: {self.email_id}") - response = datasets.sync_datasets( + response = datasets.sync_with_connection( client_id=self.client_id, project_id=self.project_id, dataset_id=self.gcs_dataset_id, @@ -169,7 +169,7 @@ def test_sync_datasets_with_multiple_data_types(self): print(f"\n Testing with data_type: {data_type}") try: - response = datasets.sync_datasets( + response = datasets.sync_with_connection( client_id=self.client_id, project_id=self.project_id, dataset_id=self.aws_dataset_id, @@ -194,7 +194,7 @@ def test_sync_datasets_invalid_connection_id(self): print("=" * 60) with self.assertRaises((LabellerrError, Exception)) as context: - datasets.sync_datasets( + datasets.sync_with_connection( client_id=self.client_id, project_id=self.project_id, dataset_id=self.aws_dataset_id, @@ -214,7 +214,7 @@ def test_sync_datasets_invalid_dataset_id(self): print("=" * 60) with self.assertRaises((LabellerrError, Exception)) as context: - datasets.sync_datasets( + datasets.sync_with_connection( client_id=self.client_id, project_id=self.project_id, dataset_id="00000000-0000-0000-0000-000000000000", From 8d3e827cbabd133cd48ca5f8f061afc698e47657 Mon Sep 17 00:00:00 2001 From: Ximi Hoque Date: Mon, 10 Nov 2025 14:55:34 +0530 Subject: [PATCH 04/13] [LABIMP-8058] create project refactored --- drivers/datasets.py | 9 +- drivers/projects.py | 65 ++ .../core/annotation_templates/__init__.py | 59 ++ labellerr/core/annotation_templates/base.py | 50 ++ labellerr/core/constants.py | 2 - labellerr/core/datasets/__init__.py | 4 +- labellerr/core/datasets/base.py | 3 - labellerr/core/datasets/datasets.py | 764 ------------------ labellerr/core/datasets/utils.py | 9 +- labellerr/core/exceptions/__init__.py | 6 + labellerr/core/projects/__init__.py | 272 +------ labellerr/core/projects/annotation_guide.py | 0 labellerr/core/projects/base.py | 8 +- labellerr/core/projects/document_project.py | 5 - .../core/schemas/annotation_templates.py | 42 + labellerr/core/schemas/datasets.py | 1 - labellerr/core/schemas/projects.py | 23 +- tests/unit/test_client.py | 365 ++++++--- tests/unit/test_create_dataset_path.py | 251 +++--- tests/unit/test_dataset_pagination.py | 54 +- 20 files changed, 617 insertions(+), 1375 deletions(-) create mode 100644 drivers/projects.py create mode 100644 labellerr/core/annotation_templates/__init__.py create mode 100644 labellerr/core/annotation_templates/base.py delete mode 100644 labellerr/core/datasets/datasets.py create mode 100644 labellerr/core/projects/annotation_guide.py create mode 100644 labellerr/core/schemas/annotation_templates.py diff --git a/drivers/datasets.py b/drivers/datasets.py index d89fcf4..365401f 100644 --- a/drivers/datasets.py +++ b/drivers/datasets.py @@ -1,7 +1,7 @@ -from labellerr.core.datasets import ( - create_dataset_from_local, - create_dataset_from_connection, -) +# from labellerr.core.datasets import ( +# create_dataset_from_local, +# create_dataset_from_connection, +# ) import logging import os from dotenv import load_dotenv @@ -58,5 +58,4 @@ dataset = LabellerrDataset( client=client, dataset_id="455e3d45-55f9-436d-98c2-07a514b7894e" ) -print(dataset.dataset_data) print(dataset.files_count) diff --git a/drivers/projects.py b/drivers/projects.py new file mode 100644 index 0000000..7c25dc5 --- /dev/null +++ b/drivers/projects.py @@ -0,0 +1,65 @@ +# from labellerr.core.projects import LabellerrProject +# from labellerr.core.annotation_templates import LabellerrAnnotationTemplate +import logging +import os +from dotenv import load_dotenv +from labellerr.client import LabellerrClient + +# from labellerr.core.schemas.annotation_templates import CreateTemplateParams, AnnotationQuestion, QuestionType +# from labellerr.core.annotation_templates import create_template +# from labellerr.core.schemas.projects import CreateProjectParams, RotationConfig +# from labellerr.core.projects import create_project +# from labellerr.core.datasets import LabellerrDataset + + +# Set logging level to DEBUG +logging.basicConfig(level=logging.DEBUG) + +load_dotenv() + +API_KEY = os.getenv("API_KEY") +API_SECRET = os.getenv("API_SECRET") +CLIENT_ID = os.getenv("CLIENT_ID") + +if not all([API_KEY, API_SECRET, CLIENT_ID]): + raise ValueError( + "API_KEY, API_SECRET, and CLIENT_ID must be set in environment variables" + ) + +# Initialize client +client = LabellerrClient( + api_key=API_KEY, + api_secret=API_SECRET, + client_id=CLIENT_ID, +) + +# project = LabellerrProject(client=client, project_id="rafaela_youngest_pike_23125") +# res = project.upload_preannotations(annotation_format="coco_json", annotation_file="/Users/Ximi-Hoque/Downloads/export_to_annotate_05_15.json") +# print(res) + +# annotation_template = LabellerrAnnotationTemplate(client=client, annotation_template_id="00016829-9051-46b1-96c6-3ec6763c342a") +# print(annotation_template.annotation_template_data) + +# res = create_template(client, CreateTemplateParams(template_name="test_template_1", data_type="image", +# questions=[ +# AnnotationQuestion(question_number=1, +# question="test_question", +# question_id="test_question_id", +# question_type=QuestionType.bounding_box, +# required=True, +# color="#FF4500", +# )])) +# print(res) + +# project = create_project(client, +# CreateProjectParams( +# project_name="test_project_via_sdk", +# data_type="image", +# rotations=RotationConfig(annotation_rotation_count=1, review_rotation_count=1, client_review_rotation_count=1), +# use_ai=False), +# datasets=[LabellerrDataset(client=client, dataset_id="ca298293-7f5e-4bdd-801f-8863a5ba458b")], +# annotation_template=res +# ) +# print (project.project_data) +# project = LabellerrProject(client=client, project_id="dinnie_confidential_lynx_20766") +# print(project.project_data) diff --git a/labellerr/core/annotation_templates/__init__.py b/labellerr/core/annotation_templates/__init__.py new file mode 100644 index 0000000..2dc8800 --- /dev/null +++ b/labellerr/core/annotation_templates/__init__.py @@ -0,0 +1,59 @@ +from .base import LabellerrAnnotationTemplate +from ..schemas.annotation_templates import CreateTemplateParams, QuestionType, Option +from .. import constants +from ..client import LabellerrClient +import uuid + +__all__ = [ + "LabellerrAnnotationTemplate", +] + +object_types = [ + QuestionType.bounding_box, + QuestionType.polygon, + QuestionType.polyline, + QuestionType.dot, +] + + +def create_template( + client: LabellerrClient, params: CreateTemplateParams +) -> LabellerrAnnotationTemplate: + """Create an annotation template""" + unique_id = str(uuid.uuid4()) + for question in params.questions: + if question.question_type in object_types and not question.color: + raise ValueError( + "Color is required for bounding box, polygon, polyline, and dot questions" + ) + + if question.question_type in object_types: + question.options = [Option(option_name=question.color)] + + # Convert questions to the expected format + questions_data = [] + for question in params.questions: + question_dict = question.model_dump() + # Convert enum to string value + question_dict["option_type"] = question.question_type.value + # Remove question_type as it's now option_type + question_dict.pop("question_type", None) + questions_data.append(question_dict) + + payload = {"templateName": params.template_name, "questions": questions_data} + url = ( + f"{constants.BASE_URL}/annotations/create_template?client_id={client.client_id}&data_type={params.data_type.value}" + f"&uuid={unique_id}" + ) + + response = client.make_request( + "POST", + url, + extra_headers={"content-type": "application/json"}, + json=payload, + request_id=unique_id, + ) + return LabellerrAnnotationTemplate( + client=client, + annotation_template_id=response.get("response", None).get("template_id"), + ) diff --git a/labellerr/core/annotation_templates/base.py b/labellerr/core/annotation_templates/base.py new file mode 100644 index 0000000..3e28aaa --- /dev/null +++ b/labellerr/core/annotation_templates/base.py @@ -0,0 +1,50 @@ +from .. import constants +from ..client import LabellerrClient +from ..exceptions import InvalidAnnotationTemplateError +import uuid + + +class LabellerrAnnotationTemplate: + @staticmethod + def get_annotation_template(client: "LabellerrClient", annotation_template_id: str): + """Get annotation template from Labellerr API""" + unique_id = str(uuid.uuid4()) + url = ( + f"{constants.BASE_URL}/annotations/get_template?template_id={annotation_template_id}&client_id={client.client_id}" + f"&uuid={unique_id}" + ) + + response = client.make_request( + "GET", + url, + extra_headers={"content-type": "application/json"}, + request_id=unique_id, + ) + return response.get("response", None) + + """Base class for all Labellerr projects with factory behavior""" + + def __new__(cls, client: "LabellerrClient", annotation_template_id: str): + # Validate that the annotation template exists before creating the instance + annotation_template_data = cls.get_annotation_template( + client, annotation_template_id + ) + + if not annotation_template_data or ( + isinstance(annotation_template_data, dict) and not annotation_template_data + ): + raise InvalidAnnotationTemplateError( + f"Annotation template with ID '{annotation_template_id}' does not exist or could not be retrieved." + ) + + # Create the instance only if validation passes + instance = super().__new__(cls) + # Store the data on the instance to avoid calling API again in __init__ + instance._annotation_template_data = annotation_template_data + return instance + + def __init__(self, client: "LabellerrClient", annotation_template_id: str): + self.client = client + self.annotation_template_id = annotation_template_id + # Use the data already fetched in __new__ + self.annotation_template_data = self._annotation_template_data diff --git a/labellerr/core/constants.py b/labellerr/core/constants.py index 6eb00b4..497e84d 100644 --- a/labellerr/core/constants.py +++ b/labellerr/core/constants.py @@ -17,8 +17,6 @@ "accepted", ] -# DATA TYPES: image, video, audio, document, text -DATA_TYPES = ("image", "video", "audio", "document", "text") DATA_TYPE_FILE_EXT = { "image": [".jpg", ".jpeg", ".png", ".tiff"], "video": [".mp4"], diff --git a/labellerr/core/datasets/__init__.py b/labellerr/core/datasets/__init__.py index c7d2707..220e89c 100644 --- a/labellerr/core/datasets/__init__.py +++ b/labellerr/core/datasets/__init__.py @@ -1,7 +1,6 @@ import json import logging import uuid -from typing import TYPE_CHECKING from .. import constants, schemas from ..exceptions import LabellerrError @@ -12,8 +11,7 @@ from .utils import upload_files, upload_folder_files_to_dataset from .video_dataset import VideoDataset as LabellerrVideoDataset -if TYPE_CHECKING: - from ..client import LabellerrClient +from ..client import LabellerrClient __all__ = [ "LabellerrImageDataset", diff --git a/labellerr/core/datasets/base.py b/labellerr/core/datasets/base.py index d980637..9360f32 100644 --- a/labellerr/core/datasets/base.py +++ b/labellerr/core/datasets/base.py @@ -6,7 +6,6 @@ from abc import ABCMeta, abstractmethod from typing import Dict, Optional, Any -from ...schemas import DataSetScope from .. import constants from ..exceptions import InvalidDatasetError from ..client import LabellerrClient @@ -52,8 +51,6 @@ def __call__(cls, client, dataset_id, **kwargs): if dataset_data is None: raise InvalidDatasetError(f"Dataset not found: {dataset_id}") data_type = dataset_data.get("data_type") - if data_type not in constants.DATA_TYPES: - raise InvalidDatasetError(f"Data type not supported: {data_type}") dataset_class = cls._registry.get(data_type) if dataset_class is None: diff --git a/labellerr/core/datasets/datasets.py b/labellerr/core/datasets/datasets.py deleted file mode 100644 index 886d9bf..0000000 --- a/labellerr/core/datasets/datasets.py +++ /dev/null @@ -1,764 +0,0 @@ -import json -import logging -import os -import uuid -from concurrent.futures import ThreadPoolExecutor, as_completed - -import requests - -from labellerr import client_utils, gcs, schemas, utils -from labellerr.core import constants -from labellerr.exceptions import LabellerrError -from labellerr.utils import validate_params - - -class DataSets(object): - """ - Handles dataset-related operations for the Labellerr API. - """ - - def __init__(self, api_key, api_secret, client): - """ - Initialize the DataSets handler. - - :param api_key: The API key for authentication - :param api_secret: The API secret for authentication - :param client: Reference to the parent Labellerr Client instance for delegating certain operations - """ - self.api_key = api_key - self.api_secret = api_secret - self.client = client - - def create_project( - self, - project_name, - data_type, - client_id, - attached_datasets, - annotation_template_id, - rotations, - use_ai=False, - created_by=None, - ): - """ - Creates a project with the given configuration. - - :param project_name: Name of the project - :param data_type: Type of data (image, video, etc.) - :param client_id: ID of the client - :param attached_datasets: List of dataset IDs to attach to the project - :param annotation_template_id: ID of the annotation template - :param rotations: Dictionary containing rotation configuration - :param use_ai: Boolean flag for AI usage (default: False) - :param created_by: Optional creator information - :return: Project creation response - :raises LabellerrError: If the creation fails - """ - # Validate parameters using Pydantic - params = schemas.CreateProjectParams( - project_name=project_name, - data_type=data_type, - client_id=client_id, - attached_datasets=attached_datasets, - annotation_template_id=annotation_template_id, - rotations=rotations, - use_ai=use_ai, - created_by=created_by, - ) - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/projects/create?client_id={params.client_id}&uuid={unique_id}" - - payload = json.dumps( - { - "project_name": params.project_name, - "attached_datasets": params.attached_datasets, - "data_type": params.data_type, - "annotation_template_id": str(params.annotation_template_id), - "rotations": params.rotations.model_dump(), - "use_ai": params.use_ai, - "created_by": params.created_by, - } - ) - - headers = client_utils.build_headers( - api_key=self.api_key, - api_secret=self.api_secret, - client_id=params.client_id, - extra_headers={ - "Origin": constants.ALLOWED_ORIGINS, - "Content-Type": "application/json", - }, - ) - - return client_utils.request( - "POST", url, headers=headers, data=payload, request_id=unique_id - ) - - def initiate_create_project(self, payload): - """ - Orchestrates project creation by handling dataset creation, annotation guidelines, - and final project setup. - """ - - try: - # validate all the parameters - required_params = [ - "client_id", - "dataset_name", - "dataset_description", - "data_type", - "created_by", - "project_name", - # Either annotation_guide or annotation_template_id must be provided - "autolabel", - ] - for param in required_params: - if param not in payload: - raise LabellerrError(f"Required parameter {param} is missing") - - if param == "client_id": - if ( - not isinstance(payload[param], str) - or not payload[param].strip() - ): - raise LabellerrError("client_id must be a non-empty string") - - # Validate created_by email format - created_by = payload.get("created_by") - if ( - not isinstance(created_by, str) - or "@" not in created_by - or "." not in created_by.split("@")[-1] - ): - raise LabellerrError("Please enter email id in created_by") - - # Ensure either annotation_guide or annotation_template_id is provided - if not payload.get("annotation_guide") and not payload.get( - "annotation_template_id" - ): - raise LabellerrError( - "Please provide either annotation guide or annotation template id" - ) - - # If annotation_guide is provided, validate its entries - if payload.get("annotation_guide"): - for guide in payload["annotation_guide"]: - if "option_type" not in guide: - raise LabellerrError( - "option_type is required in annotation_guide" - ) - if guide["option_type"] not in constants.OPTION_TYPE_LIST: - raise LabellerrError( - f"option_type must be one of {constants.OPTION_TYPE_LIST}" - ) - - if "folder_to_upload" in payload and "files_to_upload" in payload: - raise LabellerrError( - "Cannot provide both files_to_upload and folder_to_upload" - ) - - if "folder_to_upload" not in payload and "files_to_upload" not in payload: - raise LabellerrError( - "Either files_to_upload or folder_to_upload must be provided" - ) - - if ( - isinstance(payload.get("files_to_upload"), list) - and len(payload["files_to_upload"]) == 0 - ): - payload.pop("files_to_upload") - - if "rotation_config" not in payload: - payload["rotation_config"] = { - "annotation_rotation_count": 1, - "review_rotation_count": 1, - "client_review_rotation_count": 1, - } - self.validate_rotation_config(payload["rotation_config"]) - - if payload["data_type"] not in constants.DATA_TYPES: - raise LabellerrError( - f"Invalid data_type. Must be one of {constants.DATA_TYPES}" - ) - - logging.info("Rotation configuration validated . . .") - - logging.info("Creating dataset . . .") - dataset_response = self.create_dataset( - { - "client_id": payload["client_id"], - "dataset_name": payload["dataset_name"], - "data_type": payload["data_type"], - "dataset_description": payload["dataset_description"], - }, - files_to_upload=payload.get("files_to_upload"), - folder_to_upload=payload.get("folder_to_upload"), - ) - - dataset_id = dataset_response["dataset_id"] - - def check_dataset_status(): - dataset_status = self.client.get_dataset( - payload["client_id"], dataset_id - ) - - if isinstance(dataset_status, dict): - if "response" in dataset_status: - return dataset_status["response"].get("status_code", 200) == 300 - else: - return True - return False - - def is_dataset_ready(status): - return status is True - - def on_success(status): - logging.info("Dataset created and ready for use") - - def on_exception(e): - logging.error(f"Error checking dataset status: {str(e)}") - raise LabellerrError(f"Failed to check dataset status: {str(e)}") - - utils.poll( - function=check_dataset_status, - condition=is_dataset_ready, - interval=5, - on_success=on_success, - on_exception=on_exception, - ) - - if payload.get("annotation_template_id"): - annotation_template_id = payload["annotation_template_id"] - else: - annotation_template_id = self.create_annotation_guideline( - payload["client_id"], - payload["annotation_guide"], - payload["project_name"], - payload["data_type"], - ) - logging.info(f"Annotation guidelines created {annotation_template_id}") - - project_response = self.create_project( - project_name=payload["project_name"], - data_type=payload["data_type"], - client_id=payload["client_id"], - attached_datasets=[dataset_id], - annotation_template_id=annotation_template_id, - rotations=payload["rotation_config"], - use_ai=payload.get("use_ai", False), - created_by=payload["created_by"], - ) - - return { - "status": "success", - "message": "Project created successfully", - "project_id": project_response, - } - - except LabellerrError: - raise - except Exception: - logging.exception("Unexpected error in project creation") - raise - - def create_annotation_guideline( - self, client_id, questions, template_name, data_type - ): - """ - Updates the annotation guideline for a project. - - :param config: A dictionary containing the project ID, data type, client ID, autolabel status, and the annotation guideline. - :return: None - :raises LabellerrError: If the update fails. - """ - unique_id = str(uuid.uuid4()) - - url = f"{constants.BASE_URL}/annotations/create_template?data_type={data_type}&client_id={client_id}&uuid={unique_id}" - - guide_payload = json.dumps( - {"templateName": template_name, "questions": questions} - ) - - headers = client_utils.build_headers( - api_key=self.api_key, - api_secret=self.api_secret, - client_id=client_id, - extra_headers={"content-type": "application/json"}, - ) - - try: - response_data = client_utils.request( - "POST", url, headers=headers, data=guide_payload, request_id=unique_id - ) - return response_data["response"]["template_id"] - except requests.exceptions.RequestException as e: - logging.error(f"Failed to update project annotation guideline: {str(e)}") - raise - - def validate_rotation_config(self, rotation_config): - """ - Validates a rotation configuration. - - :param rotation_config: A dictionary containing the configuration for the rotations. - :raises LabellerrError: If the configuration is invalid. - """ - client_utils.validate_rotation_config(rotation_config) - - def create_dataset( - self, - dataset_config, - files_to_upload=None, - folder_to_upload=None, - connector_config=None, - ): - """ - Creates a dataset with support for multiple data types and connectors. - - :param dataset_config: A dictionary containing the configuration for the dataset. - Required fields: client_id, dataset_name, data_type - Optional fields: dataset_description, connector_type - :param files_to_upload: List of file paths to upload (for local connector) - :param folder_to_upload: Path to folder to upload (for local connector) - :param connector_config: Configuration for cloud connectors (GCP/AWS) - :return: A dictionary containing the response status and the ID of the created dataset. - """ - - try: - # Validate required fields - required_fields = ["client_id", "dataset_name", "data_type"] - for field in required_fields: - if field not in dataset_config: - raise LabellerrError( - f"Required field '{field}' missing in dataset_config" - ) - - # Validate data_type - if dataset_config.get("data_type") not in constants.DATA_TYPES: - raise LabellerrError( - f"Invalid data_type. Must be one of {constants.DATA_TYPES}" - ) - - connector_type = dataset_config.get("connector_type", "local") - connection_id = None - path = connector_type - - # Handle different connector types - if connector_type == "local": - if files_to_upload is not None: - try: - connection_id = self.client.upload_files( - client_id=dataset_config["client_id"], - files_list=files_to_upload, - ) - except Exception as e: - raise LabellerrError( - f"Failed to upload files to dataset: {str(e)}" - ) - - elif folder_to_upload is not None: - try: - result = self.upload_folder_files_to_dataset( - { - "client_id": dataset_config["client_id"], - "folder_path": folder_to_upload, - "data_type": dataset_config["data_type"], - } - ) - connection_id = result["connection_id"] - except Exception as e: - raise LabellerrError( - f"Failed to upload folder files to dataset: {str(e)}" - ) - elif connector_config is None: - # Create empty dataset for local connector - connection_id = None - - elif connector_type in ["gcp", "aws"]: - if connector_config is None: - raise LabellerrError( - f"connector_config is required for {connector_type} connector" - ) - - try: - connection_id = self.client._setup_cloud_connector( - connector_type, dataset_config["client_id"], connector_config - ) - except Exception as e: - raise LabellerrError( - f"Failed to setup {connector_type} connector: {str(e)}" - ) - else: - raise LabellerrError(f"Unsupported connector type: {connector_type}") - - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/datasets/create?client_id={dataset_config['client_id']}&uuid={unique_id}" - headers = client_utils.build_headers( - api_key=self.api_key, - api_secret=self.api_secret, - client_id=dataset_config["client_id"], - extra_headers={"content-type": "application/json"}, - ) - - payload = json.dumps( - { - "dataset_name": dataset_config["dataset_name"], - "dataset_description": dataset_config.get( - "dataset_description", "" - ), - "data_type": dataset_config["data_type"], - "connection_id": connection_id, - "path": path, - "client_id": dataset_config["client_id"], - "connector_type": connector_type, - } - ) - response_data = client_utils.request( - "POST", url, headers=headers, data=payload, request_id=unique_id - ) - dataset_id = response_data["response"]["dataset_id"] - - return {"response": "success", "dataset_id": dataset_id} - - except LabellerrError as e: - logging.error(f"Failed to create dataset: {e}") - raise - - def delete_dataset(self, client_id, dataset_id): - """ - Deletes a dataset from the system. - - :param client_id: The ID of the client - :param dataset_id: The ID of the dataset to delete - :return: Dictionary containing deletion status - :raises LabellerrError: If the deletion fails - """ - # Validate parameters using Pydantic - params = schemas.DeleteDatasetParams(client_id=client_id, dataset_id=dataset_id) - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/datasets/{params.dataset_id}/delete?client_id={params.client_id}&uuid={unique_id}" - headers = client_utils.build_headers( - api_key=self.api_key, - api_secret=self.api_secret, - client_id=params.client_id, - extra_headers={"content-type": "application/json"}, - ) - - return client_utils.request( - "DELETE", url, headers=headers, request_id=unique_id - ) - - def upload_folder_files_to_dataset(self, data_config): - """ - Uploads local files from a folder to a dataset using parallel processing. - - :param data_config: A dictionary containing the configuration for the data. - :return: A dictionary containing the response status and the list of successfully uploaded files. - :raises LabellerrError: If there are issues with file limits, permissions, or upload process - """ - try: - # Validate required fields in data_config - required_fields = ["client_id", "folder_path", "data_type"] - missing_fields = [ - field for field in required_fields if field not in data_config - ] - if missing_fields: - raise LabellerrError( - f"Missing required fields in data_config: {', '.join(missing_fields)}" - ) - - # Validate folder path exists and is accessible - if not os.path.exists(data_config["folder_path"]): - raise LabellerrError( - f"Folder path does not exist: {data_config['folder_path']}" - ) - if not os.path.isdir(data_config["folder_path"]): - raise LabellerrError( - f"Path is not a directory: {data_config['folder_path']}" - ) - if not os.access(data_config["folder_path"], os.R_OK): - raise LabellerrError( - f"No read permission for folder: {data_config['folder_path']}" - ) - - success_queue = [] - fail_queue = [] - - try: - # Get files from folder - total_file_count, total_file_volumn, filenames = ( - self.client.get_total_folder_file_count_and_total_size( - data_config["folder_path"], data_config["data_type"] - ) - ) - except Exception as e: - logging.error(f"Failed to analyze folder contents: {str(e)}") - raise - - # Check file limits - if total_file_count > constants.TOTAL_FILES_COUNT_LIMIT_PER_DATASET: - raise LabellerrError( - f"Total file count: {total_file_count} exceeds limit of {constants.TOTAL_FILES_COUNT_LIMIT_PER_DATASET} files" - ) - if total_file_volumn > constants.TOTAL_FILES_SIZE_LIMIT_PER_DATASET: - raise LabellerrError( - f"Total file size: {total_file_volumn/1024/1024:.1f}MB exceeds limit of {constants.TOTAL_FILES_SIZE_LIMIT_PER_DATASET/1024/1024:.1f}MB" - ) - - logging.info(f"Total file count: {total_file_count}") - logging.info(f"Total file size: {total_file_volumn/1024/1024:.1f} MB") - - # Use generator for memory-efficient batch creation - def create_batches(): - current_batch = [] - current_batch_size = 0 - - for file_path in filenames: - try: - file_size = os.path.getsize(file_path) - if ( - current_batch_size + file_size > constants.FILE_BATCH_SIZE - or len(current_batch) >= constants.FILE_BATCH_COUNT - ): - if current_batch: - yield current_batch - current_batch = [file_path] - current_batch_size = file_size - else: - current_batch.append(file_path) - current_batch_size += file_size - except OSError as e: - logging.error(f"Error accessing file {file_path}: {str(e)}") - fail_queue.append(file_path) - except Exception as e: - logging.error( - f"Unexpected error processing {file_path}: {str(e)}" - ) - fail_queue.append(file_path) - - if current_batch: - yield current_batch - - # Convert generator to list for ThreadPoolExecutor - batches = list(create_batches()) - - if not batches: - raise LabellerrError( - "No valid files found to upload in the specified folder" - ) - - logging.info(f"CPU count: {os.cpu_count()}, Batch Count: {len(batches)}") - - # Calculate optimal number of workers based on CPU count and batch count - max_workers = min( - os.cpu_count(), # Number of CPU cores - len(batches), # Number of batches - 20, - ) - connection_id = str(uuid.uuid4()) - # Process batches in parallel - with ThreadPoolExecutor(max_workers=max_workers) as executor: - future_to_batch = { - executor.submit( - self.__process_batch, - data_config["client_id"], - batch, - connection_id, - ): batch - for batch in batches - } - - for future in as_completed(future_to_batch): - batch = future_to_batch[future] - try: - result = future.result() - if ( - isinstance(result, dict) - and result.get("message") == "200: Success" - ): - success_queue.extend(batch) - else: - fail_queue.extend(batch) - except Exception as e: - logging.exception(e) - logging.error(f"Batch upload failed: {str(e)}") - fail_queue.extend(batch) - - if not success_queue and fail_queue: - raise LabellerrError( - "All file uploads failed. Check individual file errors above." - ) - - return { - "connection_id": connection_id, - "success": success_queue, - "fail": fail_queue, - } - - except LabellerrError: - raise - except Exception as e: - logging.error(f"Failed to upload files: {str(e)}") - raise - - def __process_batch(self, client_id, files_list, connection_id=None): - """ - Processes a batch of files. - """ - # Prepare files for upload - files = {} - for file_path in files_list: - file_name = os.path.basename(file_path) - files[file_name] = file_path - - response = self.client.connect_local_files( - client_id, list(files.keys()), connection_id - ) - resumable_upload_links = response["response"]["resumable_upload_links"] - for file_name in resumable_upload_links.keys(): - gcs.upload_to_gcs_resumable( - resumable_upload_links[file_name], files[file_name] - ) - - return response - - def attach_dataset_to_project( - self, client_id, project_id, dataset_id=None, dataset_ids=None - ): - """ - Attaches one or more datasets to an existing project. - - :param client_id: The ID of the client - :param project_id: The ID of the project - :param dataset_id: The ID of a single dataset to attach (for backward compatibility) - :param dataset_ids: List of dataset IDs to attach (for batch operations) - :return: Dictionary containing attachment status - :raises LabellerrError: If the operation fails or if neither dataset_id nor dataset_ids is provided - """ - # Handle both single and batch operations - if dataset_id is None and dataset_ids is None: - raise LabellerrError("Either dataset_id or dataset_ids must be provided") - - if dataset_id is not None and dataset_ids is not None: - raise LabellerrError( - "Cannot provide both dataset_id and dataset_ids. Use dataset_ids for batch operations." - ) - - # Convert single dataset_id to list for uniform processing - if dataset_id is not None: - dataset_ids = [dataset_id] - - # Validate parameters using Pydantic for each dataset - validated_dataset_ids = [] - for ds_id in dataset_ids: - params = schemas.AttachDatasetParams( - client_id=client_id, project_id=project_id, dataset_id=ds_id - ) - validated_dataset_ids.append(str(params.dataset_id)) - - # Use the first params validation for client_id and project_id - params = schemas.AttachDatasetParams( - client_id=client_id, project_id=project_id, dataset_id=dataset_ids[0] - ) - - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/actions/jobs/add_datasets_to_project?project_id={params.project_id}&uuid={unique_id}&client_id={params.client_id}" - headers = client_utils.build_headers( - api_key=self.api_key, - api_secret=self.api_secret, - client_id=params.client_id, - extra_headers={"content-type": "application/json"}, - ) - - payload = json.dumps({"attached_datasets": validated_dataset_ids}) - - return client_utils.request( - "POST", url, headers=headers, data=payload, request_id=unique_id - ) - - def detach_dataset_from_project( - self, client_id, project_id, dataset_id=None, dataset_ids=None - ): - """ - Detaches one or more datasets from an existing project. - - :param client_id: The ID of the client - :param project_id: The ID of the project - :param dataset_id: The ID of a single dataset to detach (for backward compatibility) - :param dataset_ids: List of dataset IDs to detach (for batch operations) - :return: Dictionary containing detachment status - :raises LabellerrError: If the operation fails or if neither dataset_id nor dataset_ids is provided - """ - # Handle both single and batch operations - if dataset_id is None and dataset_ids is None: - raise LabellerrError("Either dataset_id or dataset_ids must be provided") - - if dataset_id is not None and dataset_ids is not None: - raise LabellerrError( - "Cannot provide both dataset_id and dataset_ids. Use dataset_ids for batch operations." - ) - - # Convert single dataset_id to list for uniform processing - if dataset_id is not None: - dataset_ids = [dataset_id] - - # Validate parameters using Pydantic for each dataset - validated_dataset_ids = [] - for ds_id in dataset_ids: - params = schemas.DetachDatasetParams( - client_id=client_id, project_id=project_id, dataset_id=ds_id - ) - validated_dataset_ids.append(str(params.dataset_id)) - - # Use the first params validation for client_id and project_id - params = schemas.DetachDatasetParams( - client_id=client_id, project_id=project_id, dataset_id=dataset_ids[0] - ) - - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/actions/jobs/delete_datasets_from_project?project_id={params.project_id}&uuid={unique_id}" - headers = client_utils.build_headers( - api_key=self.api_key, - api_secret=self.api_secret, - client_id=params.client_id, - extra_headers={"content-type": "application/json"}, - ) - - payload = json.dumps({"attached_datasets": validated_dataset_ids}) - - return client_utils.request( - "POST", url, headers=headers, data=payload, request_id=unique_id - ) - - @validate_params(client_id=str, datatype=str, project_id=str, scope=str) - def get_all_datasets( - self, client_id: str, datatype: str, project_id: str, scope: str - ): - """ - Retrieves datasets by parameters. - - :param client_id: The ID of the client. - :param datatype: The type of data for the dataset. - :param project_id: The ID of the project. - :param scope: The permission scope for the dataset. - :return: The dataset list as JSON. - """ - # Validate parameters using Pydantic - params = schemas.GetAllDatasetParams( - client_id=client_id, - datatype=datatype, - project_id=project_id, - scope=scope, - ) - unique_id = str(uuid.uuid4()) - url = ( - f"{constants.BASE_URL}/datasets/list?client_id={params.client_id}&data_type={params.datatype}&permission_level={params.scope}" - f"&project_id={params.project_id}&uuid={unique_id}" - ) - headers = client_utils.build_headers( - api_key=self.api_key, - api_secret=self.api_secret, - client_id=params.client_id, - extra_headers={"content-type": "application/json"}, - ) - - return client_utils.request("GET", url, headers=headers, request_id=unique_id) diff --git a/labellerr/core/datasets/utils.py b/labellerr/core/datasets/utils.py index 2be04cb..cf1737d 100644 --- a/labellerr/core/datasets/utils.py +++ b/labellerr/core/datasets/utils.py @@ -2,14 +2,11 @@ import os import uuid from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import TYPE_CHECKING, List, Union +from typing import List, Union -from .. import client_utils, constants, gcs, schemas +from .. import client_utils, constants, gcs from ..exceptions import LabellerrError -from ..utils import validate_params - -if TYPE_CHECKING: - from ..client import LabellerrClient +from ..client import LabellerrClient def get_total_folder_file_count_and_total_size(folder_path, data_type): diff --git a/labellerr/core/exceptions/__init__.py b/labellerr/core/exceptions/__init__.py index ddb55c6..7b9e511 100644 --- a/labellerr/core/exceptions/__init__.py +++ b/labellerr/core/exceptions/__init__.py @@ -25,3 +25,9 @@ class InvalidDatasetIDError(Exception): class InvalidConnectionError(Exception): pass + + +class InvalidAnnotationTemplateError(Exception): + """Custom exception for invalid annotation template errors.""" + + pass diff --git a/labellerr/core/projects/__init__.py b/labellerr/core/projects/__init__.py index 60a41ce..212921c 100644 --- a/labellerr/core/projects/__init__.py +++ b/labellerr/core/projects/__init__.py @@ -1,256 +1,57 @@ import json -import logging import uuid -import requests - from labellerr import LabellerrClient -from .. import client_utils, constants, schemas, utils -from ..datasets import LabellerrDataset, create_dataset +from .. import client_utils, constants, schemas +from ..datasets import LabellerrDataset from ..exceptions import LabellerrError from .audio_project import AudioProject as LabellerrAudioProject -from .base import LabellerrProject from .document_project import DocucmentProject as LabellerrDocumentProject from .image_project import ImageProject as LabellerrImageProject -from .utils import validate_rotation_config from .video_project import VideoProject as LabellerrVideoProject +from .base import LabellerrProject +from ..annotation_templates import LabellerrAnnotationTemplate +from typing import List __all__ = [ - "LabellerrImageProject", - "LabellerrVideoProject", "LabellerrProject", - "LabellerrDocumentProject", "LabellerrAudioProject", + "LabellerrDocumentProject", + "LabellerrImageProject", + "LabellerrVideoProject", ] -def create_project(client: "LabellerrClient", payload: dict): +def create_project( + client: "LabellerrClient", + params: schemas.CreateProjectParams, + datasets: List[LabellerrDataset], + annotation_template: LabellerrAnnotationTemplate, +): """ Orchestrates project creation by handling dataset creation, annotation guidelines, and final project setup. """ - try: - # validate all the parameters - required_params = [ - "data_type", - "created_by", - "project_name", - # Either annotation_guide or annotation_template_id must be provided - "autolabel", - ] - for param in required_params: - if param not in payload: - raise LabellerrError(f"Required parameter {param} is missing") - - # Validate created_by email format - created_by = payload.get("created_by") - if ( - not isinstance(created_by, str) - or "@" not in created_by - or "." not in created_by.split("@")[-1] - ): - raise LabellerrError("Please enter email id in created_by") - - # Ensure either annotation_guide or annotation_template_id is provided - if not payload.get("annotation_guide") and not payload.get( - "annotation_template_id" - ): - raise LabellerrError( - "Please provide either annotation guide or annotation template id" - ) - - # If annotation_guide is provided, validate its entries - if payload.get("annotation_guide"): - for guide in payload["annotation_guide"]: - if "option_type" not in guide: - raise LabellerrError("option_type is required in annotation_guide") - if guide["option_type"] not in constants.OPTION_TYPE_LIST: - raise LabellerrError( - f"option_type must be one of {constants.OPTION_TYPE_LIST}" - ) - - if "rotation_config" not in payload: - payload["rotation_config"] = { - "annotation_rotation_count": 1, - "review_rotation_count": 1, - "client_review_rotation_count": 1, - } - validate_rotation_config(payload["rotation_config"]) - - if payload["data_type"] not in constants.DATA_TYPES: - raise LabellerrError( - f"Invalid data_type. Must be one of {constants.DATA_TYPES}" - ) - - logging.info("Rotation configuration validated . . .") - - # Handle dataset logic - either use existing datasets or create new ones - if "datasets" in payload: - # Use existing datasets - datasets = payload["datasets"] - if not isinstance(datasets, list) or len(datasets) == 0: - raise LabellerrError("datasets must be a non-empty list of dataset IDs") - - # Validate that all datasets exist and have files - logging.info("Validating existing datasets . . .") - for dataset_id in datasets: - try: - dataset = LabellerrDataset(client, dataset_id) - if dataset.files_count <= 0: - raise LabellerrError(f"Dataset {dataset_id} has no files") - except Exception as e: - raise LabellerrError( - f"Dataset {dataset_id} does not exist or is invalid: {str(e)}" - ) - - attached_datasets = datasets - logging.info("All datasets validated successfully") - else: - # Create new dataset (existing logic) - # Set dataset_name and dataset_description - if "dataset_name" not in payload: - dataset_name = payload.get("project_name") - dataset_description = ( - f"Dataset for Project - {payload.get('project_name')}" - ) - else: - dataset_name = payload.get("dataset_name") - dataset_description = payload.get( - "dataset_description", - f"Dataset for Project - {payload.get('project_name')}", - ) + if len(datasets) == 0: + raise LabellerrError("At least one dataset is required") - if "folder_to_upload" in payload and "files_to_upload" in payload: - raise LabellerrError( - "Cannot provide both files_to_upload and folder_to_upload" - ) + for dataset in datasets: + if dataset.files_count == 0: + raise LabellerrError(f"Dataset {dataset.dataset_id} has no files") - if "folder_to_upload" not in payload and "files_to_upload" not in payload: - raise LabellerrError( - "Either files_to_upload or folder_to_upload must be provided" - ) + attached_datasets = [dataset.dataset_id for dataset in datasets] - # Check for empty files_to_upload list - if ( - isinstance(payload.get("files_to_upload"), list) - and len(payload["files_to_upload"]) == 0 - ): - raise LabellerrError("files_to_upload cannot be an empty list") - - # Check for empty/whitespace folder_to_upload - if "folder_to_upload" in payload: - folder_path = payload.get("folder_to_upload", "").strip() - if not folder_path: - raise LabellerrError("Folder path does not exist") - - logging.info("Creating dataset . . .") - - dataset = create_dataset( - client, - schemas.DatasetConfig( - client_id=client.client_id, - dataset_name=dataset_name, - data_type=payload["data_type"], - dataset_description=dataset_description, - connector_type="local", - ), - files_to_upload=payload.get("files_to_upload"), - folder_to_upload=payload.get("folder_to_upload"), - ) - - def dataset_ready(): - response = LabellerrDataset( - client, dataset.dataset_id - ) # Fetch dataset again to get the status code - return response.status_code == 300 and response.files_count > 0 - - utils.poll( - function=dataset_ready, - condition=lambda x: x is True, - interval=5, - ) - - attached_datasets = [dataset.dataset_id] - logging.info("Dataset created and ready for use") - - if payload.get("annotation_template_id"): - annotation_template_id = payload["annotation_template_id"] - else: - annotation_template_id = create_annotation_guideline( - client, - payload["annotation_guide"], - payload["project_name"], - payload["data_type"], - ) - logging.info(f"Annotation guidelines created {annotation_template_id}") - project_response = __create_project_api_call( - client=client, - project_name=payload["project_name"], - data_type=payload["data_type"], - client_id=client.client_id, - attached_datasets=attached_datasets, - annotation_template_id=annotation_template_id, - rotations=payload["rotation_config"], - use_ai=payload.get("use_ai", False), - created_by=payload["created_by"], - ) - return LabellerrProject( - client, project_id=project_response["response"]["project_id"] - ) - except LabellerrError: - raise - except Exception: - logging.exception("Unexpected error in project creation") - raise - - -def __create_project_api_call( - client: "LabellerrClient", - project_name, - data_type, - client_id, - attached_datasets, - annotation_template_id, - rotations, - use_ai=False, - created_by=None, -): - """ - Creates a project with the given configuration. - - :param project_name: Name of the project - :param data_type: Type of data (image, video, etc.) - :param client_id: ID of the client - :param attached_datasets: List of dataset IDs to attach to the project - :param annotation_template_id: ID of the annotation template - :param rotations: Dictionary containing rotation configuration - :param use_ai: Boolean flag for AI usage (default: False) - :param created_by: Optional creator information - :return: Project creation response - :raises LabellerrError: If the creation fails - """ - # Validate parameters using Pydantic - params = schemas.CreateProjectParams( - project_name=project_name, - data_type=data_type, - client_id=client_id, - attached_datasets=attached_datasets, - annotation_template_id=annotation_template_id, - rotations=rotations, - use_ai=use_ai, - created_by=created_by, - ) unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/projects/create?client_id={params.client_id}&uuid={unique_id}" + url = f"{constants.BASE_URL}/projects/create?client_id={client.client_id}&uuid={unique_id}" payload = json.dumps( { "project_name": params.project_name, - "attached_datasets": params.attached_datasets, + "attached_datasets": attached_datasets, "data_type": params.data_type, - "annotation_template_id": str(params.annotation_template_id), + "annotation_template_id": annotation_template.annotation_template_id, "rotations": params.rotations.model_dump(), "use_ai": params.use_ai, "created_by": params.created_by, @@ -259,35 +60,14 @@ def __create_project_api_call( headers = client_utils.build_headers( api_key=client.api_key, api_secret=client.api_secret, - client_id=params.client_id, + client_id=client.client_id, extra_headers={ - "Origin": constants.ALLOWED_ORIGINS, "Content-Type": "application/json", }, ) - return client.make_request( + response = client.make_request( "POST", url, headers=headers, data=payload, request_id=unique_id ) - -def create_annotation_guideline( - client: "LabellerrClient", questions, template_name, data_type -): - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/annotations/create_template?data_type={data_type}&client_id={client.client_id}&uuid={unique_id}" - - guide_payload = json.dumps({"templateName": template_name, "questions": questions}) - - try: - response_data = client.make_request( - "POST", - url, - extra_headers={"content-type": "application/json"}, - request_id=unique_id, - data=guide_payload, - ) - return response_data["response"]["template_id"] - except requests.exceptions.RequestException as e: - logging.error(f"Failed to update project annotation guideline: {str(e)}") - raise + return LabellerrProject(client, project_id=response["response"]["project_id"]) diff --git a/labellerr/core/projects/annotation_guide.py b/labellerr/core/projects/annotation_guide.py new file mode 100644 index 0000000..e69de29 diff --git a/labellerr/core/projects/base.py b/labellerr/core/projects/base.py index 9af1c7a..abbd24e 100644 --- a/labellerr/core/projects/base.py +++ b/labellerr/core/projects/base.py @@ -6,7 +6,7 @@ import os import uuid from abc import ABCMeta -from typing import TYPE_CHECKING, Dict +from typing import Dict import requests @@ -15,8 +15,7 @@ from ..utils import poll from ..exports import Export -if TYPE_CHECKING: - from ..client import LabellerrClient +from ..client import LabellerrClient class LabellerrProjectMeta(ABCMeta): @@ -59,8 +58,6 @@ def __call__(cls, client, project_id, **kwargs): if project_data is None: raise InvalidProjectError(f"Project not found: {project_id}") data_type = project_data.get("data_type") - if data_type not in constants.DATA_TYPES: - raise InvalidProjectError(f"Data type not supported: {data_type}") project_class = cls._registry.get(data_type) if project_class is None: @@ -582,7 +579,6 @@ def upload_preannotations( logging.info(f"Preannotation job started successfully. Job ID: {job_id}") - # Use max_retries=10 with 5-second intervals = 50 seconds max (fits within typical test timeouts) future = self.preannotation_job_status_async(job_id) return future.result() except Exception as e: diff --git a/labellerr/core/projects/document_project.py b/labellerr/core/projects/document_project.py index 91dbb32..8885f94 100644 --- a/labellerr/core/projects/document_project.py +++ b/labellerr/core/projects/document_project.py @@ -1,11 +1,6 @@ -from typing import TYPE_CHECKING - from ..schemas import DatasetDataType from .base import LabellerrProject, LabellerrProjectMeta -if TYPE_CHECKING: - from ..client import LabellerrClient # noqa:F401 - class DocucmentProject(LabellerrProject): diff --git a/labellerr/core/schemas/annotation_templates.py b/labellerr/core/schemas/annotation_templates.py new file mode 100644 index 0000000..7016015 --- /dev/null +++ b/labellerr/core/schemas/annotation_templates.py @@ -0,0 +1,42 @@ +from pydantic import BaseModel +from typing import List, Optional +from enum import Enum +from ..schemas import DatasetDataType + + +class QuestionType(str, Enum): + bounding_box = "BoundingBox" + polygon = "polygon" + polyline = "polyline" + dot = "dot" + input = "input" + radio = "radio" + boolean = "boolean" + select = "select" + dropdown = "dropdown" + stt = "stt" + imc = "imc" + + +class Option(BaseModel): + option_name: str + + +class AnnotationQuestion(BaseModel): + """Question structure for annotation templates.""" + + question_number: int + question: str + question_id: str + question_type: QuestionType + required: bool + options: Optional[List[Option]] = [] + color: Optional[str] = None + + +class CreateTemplateParams(BaseModel): + """Parameters for creating an annotation template.""" + + template_name: str + data_type: DatasetDataType + questions: List[AnnotationQuestion] diff --git a/labellerr/core/schemas/datasets.py b/labellerr/core/schemas/datasets.py index 555fbbd..2ada187 100644 --- a/labellerr/core/schemas/datasets.py +++ b/labellerr/core/schemas/datasets.py @@ -7,7 +7,6 @@ from typing import List, Literal from uuid import UUID -from .connectors import ConnectorType from .base import DatasetDataType from pydantic import BaseModel, Field, field_validator diff --git a/labellerr/core/schemas/projects.py b/labellerr/core/schemas/projects.py index 6e01333..830853b 100644 --- a/labellerr/core/schemas/projects.py +++ b/labellerr/core/schemas/projects.py @@ -4,7 +4,9 @@ from typing import Any, Dict, List, Literal, Optional -from pydantic import BaseModel, Field, field_validator +from pydantic import BaseModel, Field + +from .base import DatasetDataType class RotationConfig(BaseModel): @@ -38,23 +40,12 @@ class CreateProjectParams(BaseModel): """Parameters for creating a project.""" project_name: str = Field(min_length=1) - data_type: Literal["image", "video", "audio", "document", "text"] - client_id: str = Field(min_length=1) - attached_datasets: List[str] = Field(min_length=1) - annotation_template_id: str + data_type: DatasetDataType rotations: RotationConfig use_ai: bool = False - created_by: Optional[str] = None - - @field_validator("attached_datasets") - @classmethod - def validate_attached_datasets(cls, v): - if not v: - raise ValueError("must contain at least one dataset ID") - for i, dataset_id in enumerate(v): - if not isinstance(dataset_id, str) or not dataset_id.strip(): - raise ValueError(f"dataset_id at index {i} must be a non-empty string") - return v + created_by: Optional[str] = Field( + None, pattern=r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$" + ) class CreateTemplateParams(BaseModel): diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 5221dc2..73f06db 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -5,8 +5,6 @@ in isolation using mocks and fixtures. """ -import os - import pytest from pydantic import ValidationError @@ -14,6 +12,10 @@ from labellerr.core.projects import create_project from labellerr.core.projects.image_project import ImageProject from labellerr.core.users.base import LabellerrUsers +from labellerr.core.schemas.projects import CreateProjectParams, RotationConfig +from labellerr.core.annotation_templates import LabellerrAnnotationTemplate +from labellerr.core.datasets import LabellerrDataset +from unittest.mock import Mock, patch @pytest.fixture @@ -41,151 +43,262 @@ def users(client): @pytest.fixture -def sample_valid_payload(): - """Create a sample valid payload for create_project""" - current_dir = os.path.dirname(os.path.abspath(__file__)) - test_image = os.path.join(current_dir, "test_data", "test_image.jpg") - - # Create test directory and file if they don't exist - os.makedirs(os.path.join(current_dir, "test_data"), exist_ok=True) - if not os.path.exists(test_image): - with open(test_image, "w") as f: - f.write("dummy image content") - - return { - "data_type": "image", - "created_by": "test_user@example.com", - "project_name": "Test Project", - "autolabel": False, - "files_to_upload": [test_image], - "annotation_guide": [ - { - "option_type": "radio", - "question": "Test Question", - "options": ["Option 1", "Option 2"], - } - ], - "rotation_config": { - "annotation_rotation_count": 1, - "review_rotation_count": 1, - "client_review_rotation_count": 1, - }, - } - - -@pytest.mark.unit -class TestInitiateCreateProject: +def sample_valid_params(): + """Create a sample valid CreateProjectParams for create_project""" + return CreateProjectParams( + project_name="Test Project", + data_type="image", + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + use_ai=False, + created_by="test_user@example.com", + ) - def test_missing_required_parameters(self, client, sample_valid_payload): - """Test error handling for missing required parameters""" - # Remove required parameters one by one and test - # Current required params in create_project: data_type, created_by, project_name, autolabel - required_params = [ - "data_type", - "created_by", - "project_name", - "autolabel", - ] - for param in required_params: - invalid_payload = sample_valid_payload.copy() - del invalid_payload[param] - - with pytest.raises(LabellerrError) as exc_info: - create_project(client, invalid_payload) +@pytest.fixture +def mock_dataset(): + """Create a mock dataset for testing""" + dataset = Mock(spec=LabellerrDataset) + dataset.dataset_id = "test-dataset-id" + dataset.files_count = 10 + return dataset - assert f"Required parameter {param} is missing" in str(exc_info.value) - # Test annotation_guide separately since it has special validation - invalid_payload = sample_valid_payload.copy() - del invalid_payload["annotation_guide"] +@pytest.fixture +def mock_annotation_template(): + """Create a mock annotation template for testing""" + template = Mock(spec=LabellerrAnnotationTemplate) + template.annotation_template_id = "test-template-id" + return template - with pytest.raises(LabellerrError) as exc_info: - create_project(client, invalid_payload) - assert ( - "Please provide either annotation guide or annotation template id" - in str(exc_info.value) - ) +@pytest.mark.unit +class TestInitiateCreateProject: - def test_invalid_created_by_email(self, client, sample_valid_payload): + def test_missing_required_parameters( + self, client, mock_dataset, mock_annotation_template + ): + """Test error handling for missing required parameters""" + # Test missing project_name + with pytest.raises(ValidationError): + CreateProjectParams( + data_type="image", + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + use_ai=False, + created_by="test_user@example.com", + # Missing project_name + ) + + # Test missing data_type + with pytest.raises(ValidationError): + CreateProjectParams( + project_name="Test Project", + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + use_ai=False, + created_by="test_user@example.com", + # Missing data_type + ) + + # Test missing rotations + with pytest.raises(ValidationError): + CreateProjectParams( + project_name="Test Project", + data_type="image", + use_ai=False, + created_by="test_user@example.com", + # Missing rotations + ) + + def test_invalid_created_by_email( + self, client, mock_dataset, mock_annotation_template + ): """Test error handling for invalid created_by email format""" - invalid_payload = sample_valid_payload.copy() - invalid_payload["created_by"] = "not_an_email" # Missing @ and domain - - with pytest.raises(LabellerrError) as exc_info: - create_project(client, invalid_payload) - - assert "Please enter email id in created_by" in str(exc_info.value) + # Test invalid email format - should raise ValidationError during CreateProjectParams creation + with pytest.raises(ValidationError): + CreateProjectParams( + project_name="Test Project", + data_type="image", + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + use_ai=False, + created_by="not_an_email", # Invalid email format + ) # Test invalid email without domain extension - invalid_payload["created_by"] = "test@example" - with pytest.raises(LabellerrError) as exc_info: - create_project(client, invalid_payload) - - assert "Please enter email id in created_by" in str(exc_info.value) - - def test_invalid_annotation_guide(self, client, sample_valid_payload): + with pytest.raises(ValidationError): + CreateProjectParams( + project_name="Test Project", + data_type="image", + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + use_ai=False, + created_by="test@example", # Invalid email format + ) + + def test_invalid_annotation_guide( + self, client, mock_dataset, mock_annotation_template + ): """Test error handling for invalid annotation guide""" - invalid_payload = sample_valid_payload.copy() - - # Missing option_type - invalid_payload["annotation_guide"] = [{"question": "Test Question"}] - with pytest.raises(LabellerrError) as exc_info: - create_project(client, invalid_payload) - - assert "option_type is required in annotation_guide" in str(exc_info.value) - - # Invalid option_type - invalid_payload["annotation_guide"] = [ - {"option_type": "invalid_type", "question": "Test Question"} - ] - with pytest.raises(LabellerrError) as exc_info: - create_project(client, invalid_payload) - - assert "option_type must be one of" in str(exc_info.value) - - def test_both_upload_methods_specified(self, client, sample_valid_payload): - """Test error when both files_to_upload and folder_to_upload are specified""" - invalid_payload = sample_valid_payload.copy() - invalid_payload["folder_to_upload"] = "/path/to/folder" + # Since annotation templates are now separate objects, + # we test that empty datasets raise an error + valid_params = CreateProjectParams( + project_name="Test Project", + data_type="image", + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + use_ai=False, + created_by="test_user@example.com", + ) + # Test with empty datasets list with pytest.raises(LabellerrError) as exc_info: - create_project(client, invalid_payload) - - assert "Cannot provide both files_to_upload and folder_to_upload" in str( - exc_info.value + create_project(client, valid_params, [], mock_annotation_template) + + assert "At least one dataset is required" in str(exc_info.value) + + def test_both_upload_methods_specified( + self, client, mock_dataset, mock_annotation_template + ): + """Test error when dataset has no files""" + valid_params = CreateProjectParams( + project_name="Test Project", + data_type="image", + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + use_ai=False, + created_by="test_user@example.com", ) - def test_no_upload_method_specified(self, client, sample_valid_payload): - """Test error when neither files_to_upload nor folder_to_upload are specified""" - invalid_payload = sample_valid_payload.copy() - del invalid_payload["files_to_upload"] + # Create a dataset with no files + empty_dataset = Mock(spec=LabellerrDataset) + empty_dataset.dataset_id = "empty-dataset-id" + empty_dataset.files_count = 0 with pytest.raises(LabellerrError) as exc_info: - create_project(client, invalid_payload) - - assert "Either files_to_upload or folder_to_upload must be provided" in str( - exc_info.value + create_project( + client, valid_params, [empty_dataset], mock_annotation_template + ) + + assert "Dataset empty-dataset-id has no files" in str(exc_info.value) + + def test_no_upload_method_specified( + self, client, mock_dataset, mock_annotation_template + ): + """Test successful project creation with valid parameters""" + valid_params = CreateProjectParams( + project_name="Test Project", + data_type="image", + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + use_ai=False, + created_by="test_user@example.com", ) - def test_empty_files_to_upload(self, client, sample_valid_payload): - """Test error handling for empty files_to_upload""" - invalid_payload = sample_valid_payload.copy() - invalid_payload["files_to_upload"] = [] - with pytest.raises(LabellerrError): - create_project(client, invalid_payload) - - def test_invalid_folder_to_upload(self, client, sample_valid_payload): - """Test error handling for invalid folder_to_upload""" - invalid_payload = sample_valid_payload.copy() - del invalid_payload["files_to_upload"] - invalid_payload["folder_to_upload"] = " " + # Mock the API response + mock_response = {"response": {"project_id": "test-project-id"}} + + with patch.object(client, "make_request", return_value=mock_response): + with patch( + "labellerr.core.projects.base.LabellerrProject.get_project", + return_value={"project_id": "test-project-id", "data_type": "image"}, + ): + result = create_project( + client, valid_params, [mock_dataset], mock_annotation_template + ) + assert result is not None + + def test_empty_files_to_upload( + self, client, mock_dataset, mock_annotation_template + ): + """Test project creation with multiple datasets""" + valid_params = CreateProjectParams( + project_name="Test Project", + data_type="image", + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + use_ai=False, + created_by="test_user@example.com", + ) - with pytest.raises(LabellerrError) as exc_info: - create_project(client, invalid_payload) + # Create multiple datasets + dataset1 = Mock(spec=LabellerrDataset) + dataset1.dataset_id = "dataset-1" + dataset1.files_count = 5 + + dataset2 = Mock(spec=LabellerrDataset) + dataset2.dataset_id = "dataset-2" + dataset2.files_count = 10 + + # Mock the API response + mock_response = {"response": {"project_id": "test-project-id"}} + + with patch.object(client, "make_request", return_value=mock_response): + with patch( + "labellerr.core.projects.base.LabellerrProject.get_project", + return_value={"project_id": "test-project-id", "data_type": "image"}, + ): + result = create_project( + client, valid_params, [dataset1, dataset2], mock_annotation_template + ) + assert result is not None + + def test_invalid_folder_to_upload( + self, client, mock_dataset, mock_annotation_template + ): + """Test project creation with AI enabled""" + valid_params = CreateProjectParams( + project_name="Test Project", + data_type="image", + rotations=RotationConfig( + annotation_rotation_count=2, + review_rotation_count=2, + client_review_rotation_count=1, + ), + use_ai=True, # Enable AI + created_by="test_user@example.com", + ) - assert "Folder path does not exist" in str(exc_info.value) + # Mock the API response + mock_response = {"response": {"project_id": "test-project-id"}} + + with patch.object(client, "make_request", return_value=mock_response): + with patch( + "labellerr.core.projects.base.LabellerrProject.get_project", + return_value={"project_id": "test-project-id", "data_type": "image"}, + ): + result = create_project( + client, valid_params, [mock_dataset], mock_annotation_template + ) + assert result is not None @pytest.mark.unit diff --git a/tests/unit/test_create_dataset_path.py b/tests/unit/test_create_dataset_path.py index 6bfa3f6..8eaa930 100644 --- a/tests/unit/test_create_dataset_path.py +++ b/tests/unit/test_create_dataset_path.py @@ -1,93 +1,31 @@ """ -Unit tests for create_dataset path parameter validation. +Unit tests for dataset creation functions. -This module focuses on testing path parameter handling for AWS and GCS connectors -in the create_dataset functionality. +This module focuses on testing the create_dataset_from_connection and +create_dataset_from_local functionality. """ from unittest.mock import patch import pytest -from labellerr.core.datasets import create_dataset +from labellerr.core.datasets import ( + create_dataset_from_connection, + create_dataset_from_local, +) from labellerr.core.exceptions import LabellerrError from labellerr.core.schemas import DatasetConfig @pytest.mark.unit -class TestCreateDatasetPathValidation: - """Test path parameter validation for AWS and GCS connectors""" +class TestCreateDatasetFunctions: + """Test dataset creation functions""" - def test_aws_connector_missing_path_with_config(self, client): - """Test that AWS connector requires path when using connector_config""" + def test_create_dataset_from_connection_success(self, client): + """Test successful dataset creation from connection""" dataset_config = DatasetConfig( client_id="test_client_id", - dataset_name="Test AWS Dataset", - data_type="image", - connector_type="aws", - ) - - aws_config = { - "aws_access_key_id": "test-key", - "aws_secret_access_key": "test-secret", - "aws_region": "us-east-1", - "bucket_name": "test-bucket", - "data_type": "image", - } - - with pytest.raises(LabellerrError) as exc_info: - create_dataset( - client=client, - dataset_config=dataset_config, - connector_config=aws_config, - # Missing path parameter - ) - - assert "path is required for aws connector" in str(exc_info.value) - - def test_gcp_connector_missing_path_with_config(self, client): - """Test that GCP connector requires path when using connector_config""" - dataset_config = DatasetConfig( - client_id="test_client_id", - dataset_name="Test GCP Dataset", - data_type="image", - connector_type="gcp", - ) - - # Create a temporary credentials file for testing - import json - import tempfile - - with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f: - json.dump({"type": "service_account"}, f) - temp_cred_file = f.name - - try: - gcp_config = { - "gcs_cred_file": temp_cred_file, - "gcs_path": "gs://test-bucket/path", - "data_type": "image", - } - - with pytest.raises(LabellerrError) as exc_info: - create_dataset( - client=client, - dataset_config=dataset_config, - connector_config=gcp_config, - # Missing path parameter - ) - - assert "path is required for gcp connector" in str(exc_info.value) - finally: - import os - - os.unlink(temp_cred_file) - - def test_aws_connector_with_path_and_connection_id(self, client): - """Test AWS connector with both path and existing connection_id""" - dataset_config = DatasetConfig( - client_id="test_client_id", - dataset_name="Test AWS Dataset", + dataset_name="Test Dataset", data_type="image", connector_type="aws", ) @@ -102,47 +40,52 @@ def test_aws_connector_with_path_and_connection_id(self, client): "labellerr.core.datasets.base.LabellerrDataset.get_dataset", return_value={"dataset_id": "test-dataset-id", "data_type": "image"}, ): - dataset = create_dataset( + dataset = create_dataset_from_connection( client=client, dataset_config=dataset_config, + connection_id="test-connection-id", path="s3://test-bucket/path/to/data", - connection_id="existing-aws-connection-id", ) - # Should succeed - path is provided with connection_id + # Should succeed assert dataset is not None - def test_gcp_connector_with_path_and_connection_id(self, client): - """Test GCP connector with both path and existing connection_id""" + def test_create_dataset_from_local_with_files(self, client): + """Test successful dataset creation from local files""" dataset_config = DatasetConfig( client_id="test_client_id", - dataset_name="Test GCP Dataset", + dataset_name="Test Local Dataset", data_type="image", - connector_type="gcp", + connector_type="local", ) mock_response = { "response": {"dataset_id": "test-dataset-id", "data_type": "image"} } - # Mock both the dataset creation and the get_dataset call - with patch.object(client, "make_request", return_value=mock_response): - with patch( - "labellerr.core.datasets.base.LabellerrDataset.get_dataset", - return_value={"dataset_id": "test-dataset-id", "data_type": "image"}, - ): - dataset = create_dataset( - client=client, - dataset_config=dataset_config, - path="gs://test-bucket/path/to/data", - connection_id="existing-gcp-connection-id", - ) - - # Should succeed - path is provided with connection_id - assert dataset is not None - - def test_local_connector_path_parameter_ignored(self, client): - """Test that local connector ignores path parameter""" + # Mock the upload_files function and dataset creation + with patch( + "labellerr.core.datasets.upload_files", return_value="test-connection-id" + ): + with patch.object(client, "make_request", return_value=mock_response): + with patch( + "labellerr.core.datasets.base.LabellerrDataset.get_dataset", + return_value={ + "dataset_id": "test-dataset-id", + "data_type": "image", + }, + ): + dataset = create_dataset_from_local( + client=client, + dataset_config=dataset_config, + files_to_upload=["test_file1.jpg", "test_file2.jpg"], + ) + + # Should succeed + assert dataset is not None + + def test_create_dataset_from_local_with_folder(self, client): + """Test successful dataset creation from local folder""" dataset_config = DatasetConfig( client_id="test_client_id", dataset_name="Test Local Dataset", @@ -154,85 +97,71 @@ def test_local_connector_path_parameter_ignored(self, client): "response": {"dataset_id": "test-dataset-id", "data_type": "image"} } - # Mock both the dataset creation and the get_dataset call - with patch.object(client, "make_request", return_value=mock_response): - with patch( - "labellerr.core.datasets.base.LabellerrDataset.get_dataset", - return_value={"dataset_id": "test-dataset-id", "data_type": "image"}, - ): - dataset = create_dataset( - client=client, - dataset_config=dataset_config, - path="some/ignored/path", # Should be ignored for local - ) - - # Should succeed - path is not validated for local connector - assert dataset is not None - - def test_aws_missing_connector_config_and_connection_id(self, client): - """Test error when neither connector_config nor connection_id is provided for AWS""" + # Mock the upload_folder_files_to_dataset function and dataset creation + with patch( + "labellerr.core.datasets.upload_folder_files_to_dataset", + return_value={"connection_id": "test-connection-id", "status": "success"}, + ): + with patch.object(client, "make_request", return_value=mock_response): + with patch( + "labellerr.core.datasets.base.LabellerrDataset.get_dataset", + return_value={ + "dataset_id": "test-dataset-id", + "data_type": "image", + }, + ): + dataset = create_dataset_from_local( + client=client, + dataset_config=dataset_config, + folder_to_upload="/path/to/test/folder", + ) + + # Should succeed + assert dataset is not None + + def test_create_dataset_from_local_no_files_or_folder(self, client): + """Test error when neither files nor folder is provided for local dataset""" dataset_config = DatasetConfig( client_id="test_client_id", - dataset_name="Test AWS Dataset", + dataset_name="Test Local Dataset", data_type="image", - connector_type="aws", + connector_type="local", ) with pytest.raises(LabellerrError) as exc_info: - create_dataset( + create_dataset_from_local( client=client, dataset_config=dataset_config, - path="s3://test-bucket/path/to/data", - # Missing both connector_config and connection_id + # Missing both files_to_upload and folder_to_upload ) - assert "connector_config is required for aws connector" in str(exc_info.value) + assert "No files or folder to upload provided" in str(exc_info.value) - def test_gcp_missing_connector_config_and_connection_id(self, client): - """Test error when neither connector_config nor connection_id is provided for GCP""" + def test_create_dataset_from_connection_with_different_paths(self, client): + """Test dataset creation with different path formats""" dataset_config = DatasetConfig( client_id="test_client_id", - dataset_name="Test GCP Dataset", + dataset_name="Test Dataset", data_type="image", connector_type="gcp", ) - with pytest.raises(LabellerrError) as exc_info: - create_dataset( - client=client, - dataset_config=dataset_config, - path="gs://test-bucket/path/to/data", - # Missing both connector_config and connection_id - ) - - assert "connector_config is required for gcp connector" in str(exc_info.value) - - def test_both_connection_id_and_connector_config(self, client): - """Test error when both connection_id and connector_config are provided""" - dataset_config = DatasetConfig( - client_id="test_client_id", - dataset_name="Test AWS Dataset", - data_type="image", - connector_type="aws", - ) - - aws_config = { - "aws_access_key_id": "test-key", - "aws_secret_access_key": "test-secret", - "aws_region": "us-east-1", - "bucket_name": "test-bucket", - "data_type": "image", + mock_response = { + "response": {"dataset_id": "test-dataset-id", "data_type": "image"} } - with pytest.raises(LabellerrError) as exc_info: - create_dataset( - client=client, - dataset_config=dataset_config, - path="s3://test-bucket/path/to/data", - connection_id="existing-connection-id", - connector_config=aws_config, - ) + # Test with GCS path + with patch.object(client, "make_request", return_value=mock_response): + with patch( + "labellerr.core.datasets.base.LabellerrDataset.get_dataset", + return_value={"dataset_id": "test-dataset-id", "data_type": "image"}, + ): + dataset = create_dataset_from_connection( + client=client, + dataset_config=dataset_config, + connection_id="test-gcp-connection-id", + path="gs://test-bucket/path/to/data", + ) - assert "Cannot provide both connection_id and connector_config" in str( - exc_info.value - ) + # Should succeed + assert dataset is not None diff --git a/tests/unit/test_dataset_pagination.py b/tests/unit/test_dataset_pagination.py index 5a0abe1..0636327 100644 --- a/tests/unit/test_dataset_pagination.py +++ b/tests/unit/test_dataset_pagination.py @@ -9,7 +9,7 @@ import pytest -from labellerr.core.datasets.base import LabellerrDataset +from labellerr.core.datasets import list_datasets from labellerr.core.schemas import DataSetScope # Helper to use correct enum values @@ -87,9 +87,7 @@ def test_default_page_size_used(self, client, mock_single_page_response): with patch.object(client, "make_request") as mock_request: mock_request.return_value = mock_single_page_response - result = LabellerrDataset.get_all_datasets( - client=client, datatype="image", scope=SCOPE_CLIENT - ) + result = list_datasets(client=client, datatype="image", scope=SCOPE_CLIENT) # Consume the generator to trigger the API call list(result) @@ -109,9 +107,7 @@ def test_default_returns_generator(self, client, mock_single_page_response): with patch.object(client, "make_request") as mock_request: mock_request.return_value = mock_single_page_response - result = LabellerrDataset.get_all_datasets( - client=client, datatype="image", scope=SCOPE_CLIENT - ) + result = list_datasets(client=client, datatype="image", scope=SCOPE_CLIENT) # Check that result is a generator import types @@ -133,7 +129,7 @@ def test_custom_page_size(self, client, mock_single_page_response): with patch.object(client, "make_request") as mock_request: mock_request.return_value = mock_single_page_response - result = LabellerrDataset.get_all_datasets( + result = list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, page_size=20 ) @@ -149,7 +145,7 @@ def test_pagination_with_last_dataset_id(self, client, mock_second_page_response with patch.object(client, "make_request") as mock_request: mock_request.return_value = mock_second_page_response - result = LabellerrDataset.get_all_datasets( + result = list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -172,7 +168,7 @@ def test_manual_pagination_flow( with patch.object(client, "make_request") as mock_request: # First page mock_request.return_value = mock_first_page_response - first_page_gen = LabellerrDataset.get_all_datasets( + first_page_gen = list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, page_size=2 ) first_page_datasets = list(first_page_gen) @@ -187,7 +183,7 @@ def test_manual_pagination_flow( # Second page mock_request.return_value = mock_last_page_response - second_page_gen = LabellerrDataset.get_all_datasets( + second_page_gen = list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -209,7 +205,7 @@ def test_auto_pagination_returns_generator(self, client, mock_single_page_respon with patch.object(client, "make_request") as mock_request: mock_request.return_value = mock_single_page_response - result = LabellerrDataset.get_all_datasets( + result = list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -229,7 +225,7 @@ def test_auto_pagination_yields_individual_datasets( mock_request.return_value = mock_single_page_response datasets = list( - LabellerrDataset.get_all_datasets( + list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -249,7 +245,7 @@ def test_auto_pagination_single_page(self, client, mock_single_page_response): mock_request.return_value = mock_single_page_response datasets = list( - LabellerrDataset.get_all_datasets( + list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -278,7 +274,7 @@ def test_auto_pagination_multiple_pages( ] datasets = list( - LabellerrDataset.get_all_datasets( + list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -305,7 +301,7 @@ def test_auto_pagination_uses_default_page_size_internally( mock_request.return_value = mock_single_page_response list( - LabellerrDataset.get_all_datasets( + list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -329,7 +325,7 @@ def test_auto_pagination_passes_last_dataset_id( ] list( - LabellerrDataset.get_all_datasets( + list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -351,7 +347,7 @@ def test_auto_pagination_early_termination( mock_second_page_response, ] - generator = LabellerrDataset.get_all_datasets( + generator = list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -385,9 +381,7 @@ def test_empty_results(self, client): with patch.object(client, "make_request") as mock_request: mock_request.return_value = empty_response - result = LabellerrDataset.get_all_datasets( - client=client, datatype="image", scope=SCOPE_CLIENT - ) + result = list_datasets(client=client, datatype="image", scope=SCOPE_CLIENT) datasets = list(result) assert datasets == [] @@ -402,7 +396,7 @@ def test_empty_results_auto_pagination(self, client): mock_request.return_value = empty_response datasets = list( - LabellerrDataset.get_all_datasets( + list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -420,7 +414,7 @@ def test_different_data_types(self, client, mock_single_page_response): with patch.object(client, "make_request") as mock_request: mock_request.return_value = mock_single_page_response - result = LabellerrDataset.get_all_datasets( + result = list_datasets( client=client, datatype=data_type, scope=SCOPE_CLIENT ) list(result) # Consume generator @@ -437,9 +431,7 @@ def test_different_scopes(self, client, mock_single_page_response): with patch.object(client, "make_request") as mock_request: mock_request.return_value = mock_single_page_response - result = LabellerrDataset.get_all_datasets( - client=client, datatype="image", scope=scope - ) + result = list_datasets(client=client, datatype="image", scope=scope) list(result) # Consume generator call_args = mock_request.call_args @@ -451,7 +443,7 @@ def test_large_page_size(self, client, mock_single_page_response): with patch.object(client, "make_request") as mock_request: mock_request.return_value = mock_single_page_response - result = LabellerrDataset.get_all_datasets( + result = list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, page_size=1000 ) list(result) # Consume generator @@ -470,7 +462,7 @@ def test_auto_pagination_memory_efficiency( mock_last_page_response, ] - generator = LabellerrDataset.get_all_datasets( + generator = list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -500,7 +492,7 @@ def test_iterate_with_for_loop( ] dataset_ids = [] - for dataset in LabellerrDataset.get_all_datasets( + for dataset in list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -522,7 +514,7 @@ def test_list_comprehension( dataset_names = [ d["name"] - for d in LabellerrDataset.get_all_datasets( + for d in list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, @@ -545,7 +537,7 @@ def test_filtering_while_iterating( # Get only datasets with even IDs even_datasets = [ d - for d in LabellerrDataset.get_all_datasets( + for d in list_datasets( client=client, datatype="image", scope=SCOPE_CLIENT, From 6684363c4701889cccf44adb85e06dd80bc6b4b0 Mon Sep 17 00:00:00 2001 From: Ximi Hoque Date: Mon, 10 Nov 2025 14:55:34 +0530 Subject: [PATCH 05/13] [LABIMP-8058] create project refactored --- labellerr/core/annotation_templates/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/labellerr/core/annotation_templates/base.py b/labellerr/core/annotation_templates/base.py index 3e28aaa..21a24af 100644 --- a/labellerr/core/annotation_templates/base.py +++ b/labellerr/core/annotation_templates/base.py @@ -40,11 +40,11 @@ def __new__(cls, client: "LabellerrClient", annotation_template_id: str): # Create the instance only if validation passes instance = super().__new__(cls) # Store the data on the instance to avoid calling API again in __init__ - instance._annotation_template_data = annotation_template_data + instance.__annotation_template_data = annotation_template_data return instance def __init__(self, client: "LabellerrClient", annotation_template_id: str): self.client = client self.annotation_template_id = annotation_template_id # Use the data already fetched in __new__ - self.annotation_template_data = self._annotation_template_data + self.annotation_template_data = self.__annotation_template_data From 8d84f523579552241619db30ee2c192daee7fda6 Mon Sep 17 00:00:00 2001 From: Ximi Hoque Date: Mon, 10 Nov 2025 15:12:36 +0530 Subject: [PATCH 06/13] List projects extracted out --- labellerr/core/projects/__init__.py | 22 ++++++++++++++++++++++ labellerr/core/projects/base.py | 23 ----------------------- 2 files changed, 22 insertions(+), 23 deletions(-) diff --git a/labellerr/core/projects/__init__.py b/labellerr/core/projects/__init__.py index 212921c..2c737be 100644 --- a/labellerr/core/projects/__init__.py +++ b/labellerr/core/projects/__init__.py @@ -71,3 +71,25 @@ def create_project( ) return LabellerrProject(client, project_id=response["response"]["project_id"]) + + +def list_projects(client: "LabellerrClient"): + """ + Retrieves a list of projects associated with a client ID. + + :param client: The client instance. + :return: A list of LabellerrProject objects. + """ + unique_id = str(uuid.uuid4()) + url = f"{constants.BASE_URL}/project_drafts/projects/detailed_list?client_id={client.client_id}&uuid={unique_id}" + + response = client.make_request( + "GET", + url, + extra_headers={"content-type": "application/json"}, + request_id=unique_id, + ) + return [ + LabellerrProject(client, project_id=project["project_id"]) + for project in response["response"]["projects"] + ] diff --git a/labellerr/core/projects/base.py b/labellerr/core/projects/base.py index abbd24e..2a1ae21 100644 --- a/labellerr/core/projects/base.py +++ b/labellerr/core/projects/base.py @@ -244,29 +244,6 @@ def update_rotation_count(self, rotation_config): logging.error(f"Project rotation update config failed: {e}") raise - @staticmethod - def list_all_projects(client: "LabellerrClient"): - """ - Retrieves a list of projects associated with a client ID. - - :param client: The client instance. - :return: A dictionary containing the list of projects. - :raises LabellerrError: If the retrieval fails. - """ - try: - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/project_drafts/projects/detailed_list?client_id={client.client_id}&uuid={unique_id}" - - return client.make_request( - "GET", - url, - extra_headers={"content-type": "application/json"}, - request_id=unique_id, - ) - except Exception as e: - logging.error(f"Failed to retrieve projects: {str(e)}") - raise - def _upload_preannotation_sync( self, project_id, From 10c43239ebec132e21d3499e863578238e0b4053 Mon Sep 17 00:00:00 2001 From: Ximi Hoque Date: Mon, 10 Nov 2025 19:16:19 +0530 Subject: [PATCH 07/13] Updated few more functions and preannotations, unit tests --- driver.py | 4 +- drivers/projects.py | 2 +- .../core/annotation_templates/__init__.py | 15 +- labellerr/core/connectors/connections.py | 9 +- labellerr/core/datasets/__init__.py | 5 +- labellerr/core/projects/base.py | 286 +++--------------- labellerr/core/projects/utils.py | 38 ++- .../integration/test_labellerr_integration.py | 15 +- tests/unit/test_client.py | 2 +- tests/unit/test_create_dataset_path.py | 51 +++- tests/unit/test_keyframes.py | 2 +- 11 files changed, 153 insertions(+), 276 deletions(-) diff --git a/driver.py b/driver.py index bc71642..ac37e32 100644 --- a/driver.py +++ b/driver.py @@ -181,9 +181,9 @@ # print(file.file_data) # project = LabellerrProject(client=client, project_id="") -# res = project.upload_preannotations( +# res = project.upload_preannotation( # annotation_format="coco_json", annotation_file="horses_coco.json" -# ) +# ).result() # print(res) # print(LabellerrProject.list_all_projects(client=client)) diff --git a/drivers/projects.py b/drivers/projects.py index 7c25dc5..8e07c39 100644 --- a/drivers/projects.py +++ b/drivers/projects.py @@ -34,7 +34,7 @@ ) # project = LabellerrProject(client=client, project_id="rafaela_youngest_pike_23125") -# res = project.upload_preannotations(annotation_format="coco_json", annotation_file="/Users/Ximi-Hoque/Downloads/export_to_annotate_05_15.json") +# res = project.upload_preannotation(annotation_format="coco_json", annotation_file="/Users/Ximi-Hoque/Downloads/export_to_annotate_05_15.json").result() # print(res) # annotation_template = LabellerrAnnotationTemplate(client=client, annotation_template_id="00016829-9051-46b1-96c6-3ec6763c342a") diff --git a/labellerr/core/annotation_templates/__init__.py b/labellerr/core/annotation_templates/__init__.py index 2dc8800..f89441b 100644 --- a/labellerr/core/annotation_templates/__init__.py +++ b/labellerr/core/annotation_templates/__init__.py @@ -20,15 +20,20 @@ def create_template( client: LabellerrClient, params: CreateTemplateParams ) -> LabellerrAnnotationTemplate: """Create an annotation template""" + unique_id = str(uuid.uuid4()) for question in params.questions: - if question.question_type in object_types and not question.color: - raise ValueError( - "Color is required for bounding box, polygon, polyline, and dot questions" - ) - if question.question_type in object_types: + if not question.color: + raise ValueError( + "Color is required for bounding box, polygon, polyline, and dot questions" + ) question.options = [Option(option_name=question.color)] + else: + if question.question_type != QuestionType.input and not question.options: + raise ValueError( + "Options are required for radio, boolean, select, dropdown, stt, imc questions" + ) # Convert questions to the expected format questions_data = [] diff --git a/labellerr/core/connectors/connections.py b/labellerr/core/connectors/connections.py index c691813..0c56f27 100644 --- a/labellerr/core/connectors/connections.py +++ b/labellerr/core/connectors/connections.py @@ -2,13 +2,12 @@ import uuid from abc import ABCMeta -from typing import TYPE_CHECKING, Dict +from typing import Dict from .. import client_utils, constants from ..schemas import ConnectionType, DatasetDataType -from ..exceptions import InvalidConnectionError, InvalidDatasetIDError +from ..exceptions import InvalidConnectionError -if TYPE_CHECKING: - from ..client import LabellerrClient +from ..client import LabellerrClient class LabellerrConnectionMeta(ABCMeta): @@ -49,7 +48,7 @@ def __call__(cls, client, connection_id, **kwargs): return instance connection_data = cls.get_connection(client, connection_id) if connection_data is None: - raise InvalidDatasetIDError(f"Connection not found: {connection_id}") + raise InvalidConnectionError(f"Connection not found: {connection_id}") connector = connection_data.get("connector") connection_class = cls._registry.get(connector) if connection_class is None: diff --git a/labellerr/core/datasets/__init__.py b/labellerr/core/datasets/__init__.py index 220e89c..46b1b18 100644 --- a/labellerr/core/datasets/__init__.py +++ b/labellerr/core/datasets/__init__.py @@ -10,6 +10,7 @@ from .image_dataset import ImageDataset as LabellerrImageDataset from .utils import upload_files, upload_folder_files_to_dataset from .video_dataset import VideoDataset as LabellerrVideoDataset +from ..connectors import LabellerrConnection from ..client import LabellerrClient @@ -25,7 +26,7 @@ def create_dataset_from_connection( client: "LabellerrClient", dataset_config: schemas.DatasetConfig, - connection_id: str, + connection: LabellerrConnection, path: str, ) -> LabellerrDataset: """ @@ -45,7 +46,7 @@ def create_dataset_from_connection( "dataset_name": dataset_config.dataset_name, "dataset_description": dataset_config.dataset_description, "data_type": dataset_config.data_type, - "connection_id": connection_id, + "connection_id": connection.connection_id, "path": path, "client_id": client.client_id, "es_multimodal_index": dataset_config.multimodal_indexing, diff --git a/labellerr/core/projects/base.py b/labellerr/core/projects/base.py index 2a1ae21..14e05de 100644 --- a/labellerr/core/projects/base.py +++ b/labellerr/core/projects/base.py @@ -10,9 +10,9 @@ import requests -from .. import client_utils, constants, gcs, schemas +from .. import client_utils, constants, schemas from ..exceptions import InvalidProjectError, LabellerrError -from ..utils import poll +from .utils import poll from ..exports import Export from ..client import LabellerrClient @@ -72,44 +72,31 @@ class LabellerrProject(metaclass=LabellerrProjectMeta): def __init__(self, client: "LabellerrClient", project_id: str, **kwargs): self.client = client self.project_id = project_id - self.project_data = kwargs["project_data"] + self.__project_data = kwargs["project_data"] @property - def data_type(self): - return self.project_data.get("data_type") + def status_code(self): + return self.__project_data.get("status_code", 501) # if not found, return 501 @property - def attached_datasets(self): - return self.project_data.get("attached_datasets") + def annotation_template_id(self): + return self.__project_data.get("annotation_template_id") - def get_direct_upload_url( - self, file_name: str, client_id: str, purpose: str = "pre-annotations" - ) -> str: - """ - Get a direct upload URL for uploading files to GCS. + @property + def created_by(self): + return self.__project_data.get("created_by") - :param file_name: Name of the file to upload - :param client_id: Client ID - :param purpose: Purpose of the upload (default: "pre-annotations") - :return: Direct upload URL - """ - url = f"{constants.BASE_URL}/connectors/direct-upload-url" - params = { # noqa: F841 - "client_id": client_id, - "purpose": purpose, - "file_name": file_name, - } + @property + def created_at(self): + return self.__project_data.get("created_at") - try: - response_data = self.client.make_request( - "GET", - url, - extra_headers={"Origin": constants.ALLOWED_ORIGINS}, - ) - return response_data["response"] - except Exception as e: - logging.error(f"Error getting direct upload url: {e}") - raise LabellerrError(f"Failed to get direct upload URL: {str(e)}") + @property + def data_type(self): + return self.__project_data.get("data_type") + + @property + def attached_datasets(self): + return self.__project_data.get("attached_datasets") def detach_dataset_from_project(self, dataset_id=None, dataset_ids=None): """ @@ -244,204 +231,6 @@ def update_rotation_count(self, rotation_config): logging.error(f"Project rotation update config failed: {e}") raise - def _upload_preannotation_sync( - self, - project_id, - client_id, - annotation_format, - annotation_file, - conf_bucket=None, - ): - """ - Synchronous implementation of preannotation upload. - - :param project_id: The ID of the project. - :param client_id: The ID of the client. - :param annotation_format: The format of the preannotation data. - :param annotation_file: The file path of the preannotation data. - :param conf_bucket: Confidence bucket [low, medium, high] - :return: The response from the API. - :raises LabellerrError: If the upload fails. - """ - try: - # validate all the parameters - required_params = { - "project_id": project_id, - "client_id": client_id, - "annotation_format": annotation_format, - "annotation_file": annotation_file, - } - client_utils.validate_required_params( - required_params, list(required_params.keys()) - ) - client_utils.validate_annotation_format(annotation_format, annotation_file) - - request_uuid = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/actions/upload_answers?project_id={project_id}&answer_format={annotation_format}&client_id={client_id}&uuid={request_uuid}" - if conf_bucket: - assert conf_bucket in [ - "low", - "medium", - "high", - ], "Invalid confidence bucket value. Must be one of [low, medium, high]" - url += f"&conf_bucket={conf_bucket}" - file_name = client_utils.validate_file_exists(annotation_file) - # get the direct upload url - gcs_path = f"{project_id}/{annotation_format}-{file_name}" - logging.info("Uploading your file to Labellerr. Please wait...") - direct_upload_url = self.get_direct_upload_url(gcs_path, client_id) - # Now let's wait for the file to be uploaded to the gcs - gcs.upload_to_gcs_direct(direct_upload_url, annotation_file) - payload = {} - url += "&gcs_path=" + gcs_path - - response = self.client.make_request( - "POST", - url, - extra_headers={"email_id": self.client.api_key}, - request_id=request_uuid, - handle_response=False, - data=payload, - ) - response_data = self.client.handle_upload_response(response, request_uuid) - - # read job_id from the response - job_id = response_data["response"]["job_id"] - self.client_id = client_id - - logging.info(f"Preannotation upload successful. Job ID: {job_id}") - - # Use max_retries=10 with 5-second intervals = 50 seconds max (fits within typical test timeouts) - future = self.preannotation_job_status_async(project_id, job_id) - return future.result() - except Exception as e: - logging.error(f"Failed to upload preannotation: {str(e)}") - raise - - def upload_preannotation_async( - self, annotation_format, annotation_file, conf_bucket=None - ): - """ - Asynchronously uploads preannotation data to a project. - - :param annotation_format: The format of the preannotation data. - :param annotation_file: The file path of the preannotation data. - :param conf_bucket: Confidence bucket [low, medium, high] - :return: A Future object that will contain the response from the API. - :raises LabellerrError: If the upload fails. - - """ - - def upload_and_monitor(): - try: - # validate all the parameters - required_params = [ - "annotation_format", - "annotation_file", - ] - for param in required_params: - if param not in locals(): - raise LabellerrError(f"Required parameter {param} is missing") - - if annotation_format not in constants.ANNOTATION_FORMAT: - raise LabellerrError( - f"Invalid annotation_format. Must be one of {constants.ANNOTATION_FORMAT}" - ) - - request_uuid = str(uuid.uuid4()) - url = ( - f"{constants.BASE_URL}/actions/upload_answers?" - f"project_id={self.project_id}&answer_format={annotation_format}&client_id={self.client.client_id}&uuid={request_uuid}" - ) - if conf_bucket: - assert conf_bucket in [ - "low", - "medium", - "high", - ], "Invalid confidence bucket value. Must be one of [low, medium, high]" - url += f"&conf_bucket={conf_bucket}" - # validate if the file exist then extract file name from the path - if os.path.exists(annotation_file): - file_name = os.path.basename(annotation_file) - else: - raise LabellerrError("File not found") - - # Check if the file extension is .json when annotation_format is coco_json - if annotation_format == "coco_json": - file_extension = os.path.splitext(annotation_file)[1].lower() - if file_extension != ".json": - raise LabellerrError( - "For coco_json annotation format, the file must have a .json extension" - ) - # get the direct upload url - gcs_path = f"{self.project_id}/{annotation_format}-{file_name}" - logging.info("Uploading your file to Labellerr. Please wait...") - direct_upload_url = self.get_direct_upload_url( - gcs_path, self.client.client_id - ) - # Now let's wait for the file to be uploaded to the gcs - gcs.upload_to_gcs_direct(direct_upload_url, annotation_file) - payload = {} - url += "&gcs_path=" + gcs_path - - response = self.client.make_request( - "POST", - url, - extra_headers={"email_id": self.client.api_key}, - request_id=request_uuid, - handle_response=False, - data=payload, - ) - response_data = self.client.handle_upload_response( - response, request_uuid - ) - - # read job_id from the response - job_id = response_data["response"]["job_id"] - - logging.info(f"Pre annotation upload successful. Job ID: {job_id}") - - # Now monitor the status - status_url = f"{constants.BASE_URL}/actions/upload_answers_status?project_id={self.project_id}&job_id={job_id}&client_id={self.client.client_id}" - - def check_job_status(): - status_data = self.client.make_request( - "GET", - status_url, - extra_headers={"Origin": constants.ALLOWED_ORIGINS}, - ) - logging.debug(f"Status data: {status_data}") - return status_data - - def is_job_completed(status_data): - return status_data.get("response", {}).get("status") == "completed" - - def on_success(status_data): - logging.info("Pre-annotation job completed.") - - def on_exception(e): - logging.error(f"Failed to get preannotation job status: {str(e)}") - raise LabellerrError( - f"Failed to get preannotation job status: {str(e)}" - ) - - result = poll( - function=check_job_status, - condition=is_job_completed, - interval=5.0, - on_success=on_success, - on_exception=on_exception, - ) - - return result - - except Exception as e: - logging.exception(f"Failed to upload preannotation: {str(e)}") - raise - - with concurrent.futures.ThreadPoolExecutor() as executor: - return executor.submit(upload_and_monitor) - def preannotation_job_status_async(self, job_id): """ Get the status of a preannotation job asynchronously with timeout protection. @@ -495,16 +284,19 @@ def on_exception(e): with concurrent.futures.ThreadPoolExecutor() as executor: return executor.submit(check_status) - def upload_preannotations( - self, annotation_format, annotation_file, conf_bucket=None + def __upload_preannotations( + self, + annotation_format, + annotation_file, + conf_bucket=None, ): """ - Uploads preannotation data to a project. + Uploads preannotation data to a project asynchronously. :param annotation_format: The format of the preannotation data. :param annotation_file: The file path of the preannotation data. :param conf_bucket: Confidence bucket [low, medium, high] - :return: The response from the API. + :return: concurrent.futures.Future - call .result() to wait for completion :raises LabellerrError: If the upload fails. """ try: @@ -556,12 +348,32 @@ def upload_preannotations( logging.info(f"Preannotation job started successfully. Job ID: {job_id}") - future = self.preannotation_job_status_async(job_id) - return future.result() + return self.preannotation_job_status_async(job_id) except Exception as e: logging.error(f"Failed to upload preannotation: {str(e)}") raise + def upload_preannotations( + self, annotation_format, annotation_file, conf_bucket=None, _async=False + ): + """ + Backward compatibility method for upload_preannotations. + + :param annotation_format: The format of the preannotation data. + :param annotation_file: The file path of the preannotation data. + :param conf_bucket: Confidence bucket [low, medium, high] + :param _async: Whether to return a future object (True) or block until completion (False) + :return: The response from the API (blocks until completion). + :raises LabellerrError: If the upload fails. + """ + future = self.__upload_preannotations( + annotation_format, annotation_file, conf_bucket + ) + if _async: + return future + else: + return future.result() + def create_export(self, export_config: schemas.CreateExportParams): """ Creates an export with the given configuration. diff --git a/labellerr/core/projects/utils.py b/labellerr/core/projects/utils.py index 0060a3d..4490afa 100644 --- a/labellerr/core/projects/utils.py +++ b/labellerr/core/projects/utils.py @@ -1,7 +1,14 @@ from typing import Any, Dict from ..exceptions import LabellerrError -from ..utils import poll # noqa: F401 +from ..utils import poll + +from .. import constants +from ..client import LabellerrClient + +__all__ = [ + "poll", +] def validate_rotation_config(rotation_config: Dict[str, Any]) -> None: @@ -40,3 +47,32 @@ def validate_rotation_config(rotation_config: Dict[str, Any]) -> None: raise LabellerrError( "client_review_rotation_count must be 0 when annotation_rotation_count is greater than 1" ) + + +def get_direct_upload_url( + client: LabellerrClient, file_name: str, purpose: str = "pre-annotations" +) -> str: + """ + Get a direct upload URL for uploading files to GCS. + + :param file_name: Name of the file to upload + :param client: LabellerrClient instance + :param purpose: Purpose of the upload (default: "pre-annotations") + :return: Direct upload URL + """ + url = f"{constants.BASE_URL}/connectors/direct-upload-url" + params = { # noqa: F841 + "client_id": client.client_id, + "purpose": purpose, + "file_name": file_name, + } + + try: + response_data = client.make_request( + "GET", + url, + extra_headers={"Origin": constants.ALLOWED_ORIGINS}, + ) + return response_data["response"] + except Exception as e: + raise LabellerrError(f"Failed to get direct upload URL: {str(e)}") diff --git a/tests/integration/test_labellerr_integration.py b/tests/integration/test_labellerr_integration.py index 9b12513..2629803 100644 --- a/tests/integration/test_labellerr_integration.py +++ b/tests/integration/test_labellerr_integration.py @@ -137,12 +137,11 @@ def test_pre_annotation_upload_coco_json( annotation_file = temp_json_file(sample_annotation_data["coco_json"]) try: - result = project._upload_preannotation_sync( - project_id=test_project_ids["project_id"], - client_id=test_credentials["client_id"], + future = project.upload_preannotation( annotation_format="coco_json", annotation_file=annotation_file, ) + result = future.result() assert isinstance(result, dict) assert "response" in result @@ -183,12 +182,11 @@ def timeout_handler(signum, frame): signal.alarm(60) try: - result = project._upload_preannotation_sync( - project_id=test_project_ids["project_id"], - client_id=test_credentials["client_id"], + future = project.upload_preannotation( annotation_format="json", annotation_file=annotation_file, ) + result = future.result() assert isinstance(result, dict) @@ -230,12 +228,11 @@ def test_pre_annotation_invalid_format( project = LabellerrProject(integration_client, test_project_ids["project_id"]) with pytest.raises(LabellerrError) as exc_info: - project._upload_preannotation_sync( - project_id=test_project_ids["project_id"], - client_id=test_credentials["client_id"], + future = project.upload_preannotation( annotation_format=invalid_format, annotation_file="test.json", ) + future.result() assert expected_error in str(exc_info.value) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 73f06db..96df217 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -31,7 +31,7 @@ def project(client): proj = ImageProject.__new__(ImageProject) proj.client = client proj.project_id = "test_project_id" - proj.project_data = project_data + proj.__project_data = project_data return proj diff --git a/tests/unit/test_create_dataset_path.py b/tests/unit/test_create_dataset_path.py index 8eaa930..eee98a7 100644 --- a/tests/unit/test_create_dataset_path.py +++ b/tests/unit/test_create_dataset_path.py @@ -13,6 +13,7 @@ create_dataset_from_connection, create_dataset_from_local, ) +from labellerr.core.datasets.base import LabellerrDataset from labellerr.core.exceptions import LabellerrError from labellerr.core.schemas import DatasetConfig @@ -40,10 +41,16 @@ def test_create_dataset_from_connection_success(self, client): "labellerr.core.datasets.base.LabellerrDataset.get_dataset", return_value={"dataset_id": "test-dataset-id", "data_type": "image"}, ): + # Create a mock connection object + from unittest.mock import Mock + + mock_connection = Mock() + mock_connection.connection_id = "test-connection-id" + dataset = create_dataset_from_connection( client=client, dataset_config=dataset_config, - connection_id="test-connection-id", + connection=mock_connection, path="s3://test-bucket/path/to/data", ) @@ -75,11 +82,18 @@ def test_create_dataset_from_local_with_files(self, client): "data_type": "image", }, ): - dataset = create_dataset_from_local( - client=client, - dataset_config=dataset_config, - files_to_upload=["test_file1.jpg", "test_file2.jpg"], - ) + # Mock create_dataset_from_connection to avoid the connection object issue + with patch( + "labellerr.core.datasets.create_dataset_from_connection", + return_value=LabellerrDataset( + client=client, dataset_id="test-dataset-id" + ), + ): + dataset = create_dataset_from_local( + client=client, + dataset_config=dataset_config, + files_to_upload=["test_file1.jpg", "test_file2.jpg"], + ) # Should succeed assert dataset is not None @@ -110,11 +124,18 @@ def test_create_dataset_from_local_with_folder(self, client): "data_type": "image", }, ): - dataset = create_dataset_from_local( - client=client, - dataset_config=dataset_config, - folder_to_upload="/path/to/test/folder", - ) + # Mock create_dataset_from_connection to avoid the connection object issue + with patch( + "labellerr.core.datasets.create_dataset_from_connection", + return_value=LabellerrDataset( + client=client, dataset_id="test-dataset-id" + ), + ): + dataset = create_dataset_from_local( + client=client, + dataset_config=dataset_config, + folder_to_upload="/path/to/test/folder", + ) # Should succeed assert dataset is not None @@ -156,10 +177,16 @@ def test_create_dataset_from_connection_with_different_paths(self, client): "labellerr.core.datasets.base.LabellerrDataset.get_dataset", return_value={"dataset_id": "test-dataset-id", "data_type": "image"}, ): + # Create a mock connection object + from unittest.mock import Mock + + mock_connection = Mock() + mock_connection.connection_id = "test-gcp-connection-id" + dataset = create_dataset_from_connection( client=client, dataset_config=dataset_config, - connection_id="test-gcp-connection-id", + connection=mock_connection, path="gs://test-bucket/path/to/data", ) diff --git a/tests/unit/test_keyframes.py b/tests/unit/test_keyframes.py index 132fc42..9650254 100644 --- a/tests/unit/test_keyframes.py +++ b/tests/unit/test_keyframes.py @@ -201,7 +201,7 @@ def mock_video_project(mock_client): project = VideoProject.__new__(VideoProject) project.client = mock_client project.project_id = "test_project_id" - project.project_data = { + project.__project_data = { "project_id": "test_project_id", "data_type": "video", "attached_datasets": [], From ac4d5de7d0d6df3b5642959a832fcb9ec56dfe8d Mon Sep 17 00:00:00 2001 From: Ximi Hoque Date: Tue, 11 Nov 2025 19:39:02 +0530 Subject: [PATCH 08/13] [LABIMP-8041] Labellerr File class updates, dataset updates --- drivers/datasets.py | 3 + drivers/projects.py | 2 +- labellerr/core/datasets/base.py | 75 +++++++++- labellerr/core/datasets/image_dataset.py | 3 +- labellerr/core/datasets/video_dataset.py | 86 ----------- labellerr/core/files/base.py | 11 +- labellerr/core/files/video_file.py | 4 +- labellerr/core/projects/video_project.py | 78 ++++------ .../core/schemas/annotation_templates.py | 6 +- labellerr/core/utils.py | 140 ------------------ 10 files changed, 115 insertions(+), 293 deletions(-) delete mode 100644 labellerr/core/utils.py diff --git a/drivers/datasets.py b/drivers/datasets.py index 365401f..f57356e 100644 --- a/drivers/datasets.py +++ b/drivers/datasets.py @@ -59,3 +59,6 @@ client=client, dataset_id="455e3d45-55f9-436d-98c2-07a514b7894e" ) print(dataset.files_count) + +for file in dataset.fetch_files(): + print(file.file_id) diff --git a/drivers/projects.py b/drivers/projects.py index 8e07c39..1afb013 100644 --- a/drivers/projects.py +++ b/drivers/projects.py @@ -8,7 +8,7 @@ # from labellerr.core.schemas.annotation_templates import CreateTemplateParams, AnnotationQuestion, QuestionType # from labellerr.core.annotation_templates import create_template # from labellerr.core.schemas.projects import CreateProjectParams, RotationConfig -# from labellerr.core.projects import create_project +# from labellerr.core.projects import create_project, list_projects # from labellerr.core.datasets import LabellerrDataset diff --git a/labellerr/core/datasets/base.py b/labellerr/core/datasets/base.py index 9360f32..be01bfa 100644 --- a/labellerr/core/datasets/base.py +++ b/labellerr/core/datasets/base.py @@ -3,13 +3,15 @@ import json import logging import uuid -from abc import ABCMeta, abstractmethod -from typing import Dict, Optional, Any +from abc import ABCMeta +from typing import Dict, Optional, Any, List from .. import constants -from ..exceptions import InvalidDatasetError +from ..exceptions import InvalidDatasetError, LabellerrError from ..client import LabellerrClient +from ..files import LabellerrFile + class LabellerrDatasetMeta(ABCMeta): # Class-level registry for dataset types @@ -170,10 +172,69 @@ def on_success(dataset_data): on_success=on_success, ) - @abstractmethod - def fetch_files(self): - """Each file type must implement its own download logic""" - pass + def fetch_files(self, page_size: int = 1000) -> List[LabellerrFile]: + """ + Fetch all files in this dataset as LabellerrFile instances. + + :param page_size: Number of files to fetch per API request (default: 10) + :return: List of file IDs + """ + print(f"Fetching files for dataset: {self.dataset_id}") + file_ids = [] + next_search_after = None # Start with None for first page + + while True: + unique_id = str(uuid.uuid4()) + url = f"{constants.BASE_URL}/search/files/all" + params = { + "sort_by": "created_at", + "sort_order": "desc", + "size": page_size, + "uuid": unique_id, + "dataset_id": self.dataset_id, + "client_id": self.client.client_id, + } + + # Add next_search_after only if it exists (don't send on first request) + if next_search_after: + url += f"?next_search_after={next_search_after}" + + response = self.client.make_request( + "GET", url, extra_headers=None, request_id=unique_id, params=params + ) + print(response) + # Extract files from the response + files = response.get("response", {}).get("files", []) + + # Collect file IDs + for file_info in files: + file_id = file_info.get("file_id") + if file_id: + file_ids.append(file_id) + + # Get next_search_after for pagination + next_search_after = response.get("response", {}).get("next_search_after") + + # Break if no more pages or no files returned + if not next_search_after or not files: + break + + files = [] + + for file_id in file_ids: + try: + _file = LabellerrFile( + client=self.client, + file_id=file_id, + dataset_id=self.dataset_id, + ) + files.append(_file) + except LabellerrError as e: + logging.warning( + f"Warning: Failed to create file instance for {file_id}: {str(e)}" + ) + + return files def sync_with_connection( self, diff --git a/labellerr/core/datasets/image_dataset.py b/labellerr/core/datasets/image_dataset.py index 9c46ec5..c08edfa 100644 --- a/labellerr/core/datasets/image_dataset.py +++ b/labellerr/core/datasets/image_dataset.py @@ -3,8 +3,7 @@ class ImageDataset(LabellerrDataset): - def fetch_files(self): - print("Yo I am gonna fetch some files!") + pass LabellerrDatasetMeta._register(DatasetDataType.image, ImageDataset) diff --git a/labellerr/core/datasets/video_dataset.py b/labellerr/core/datasets/video_dataset.py index 6f38301..6e8538e 100644 --- a/labellerr/core/datasets/video_dataset.py +++ b/labellerr/core/datasets/video_dataset.py @@ -1,8 +1,4 @@ -import uuid - -from .. import constants from ..exceptions import LabellerrError -from ..files import LabellerrFile from ..schemas import DatasetDataType from .base import LabellerrDataset, LabellerrDatasetMeta @@ -12,88 +8,6 @@ class VideoDataset(LabellerrDataset): Class for handling video dataset operations and fetching multiple video files. """ - def fetch_files(self, page_size: int = 1000): - """ - Fetch all video files in this dataset as LabellerrVideoFile instances. - - :param page_size: Number of files to fetch per API request (default: 10) - :return: List of file IDs - """ - try: - all_file_ids = [] - next_search_after = None # Start with None for first page - - while True: - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/search/files/all" - params = { - "sort_by": "created_at", - "sort_order": "desc", - "size": page_size, - "uuid": unique_id, - "dataset_id": self.dataset_id, - "client_id": self.client.client_id, - } - - # Add next_search_after only if it exists (don't send on first request) - if next_search_after: - url += f"?next_search_after={next_search_after}" - - # print(params) - - response = self.client.make_request(url, params, unique_id) - - # pprint.pprint(response) - - # Extract files from the response - files = response.get("response", {}).get("files", []) - - # Collect file IDs - for file_info in files: - file_id = file_info.get("file_id") - if file_id: - all_file_ids.append(file_id) - - # Get next_search_after for pagination - next_search_after = response.get("response", {}).get( - "next_search_after" - ) - - # Break if no more pages or no files returned - if not next_search_after or not files: - break - - print(f"Fetched total: {len(all_file_ids)}") - - print(f"Total file IDs extracted: {len(all_file_ids)}") - # return all_file_ids - - # Create LabellerrVideoFile instances for each file_id - video_files = [] - print( - f"\nCreating LabellerrFile instances for {len(all_file_ids)} files..." - ) - - for file_id in all_file_ids: - try: - video_file = LabellerrFile( - client=self.client, - file_id=file_id, - project_id="self.project_id", # noqa: # todo: ximi we don't have project id here - dataset_id=self.dataset_id, - ) - video_files.append(video_file) - except Exception as e: - print( - f"Warning: Failed to create file instance for {file_id}: {str(e)}" - ) - - print(f"Successfully created {len(video_files)} LabellerrFile instances") - return video_files - - except Exception as e: - raise LabellerrError(f"Failed to fetch dataset files: {str(e)}") - def download(self): """ Process all video files in the dataset: download frames, create videos, diff --git a/labellerr/core/files/base.py b/labellerr/core/files/base.py index 98bcf09..b504e1d 100644 --- a/labellerr/core/files/base.py +++ b/labellerr/core/files/base.py @@ -51,7 +51,6 @@ def __call__( elif dataset_id: params["dataset_id"] = dataset_id - # TODO: Add dataset_id to params based on precedence logic # Priority: project_id > dataset_id url = f"{constants.BASE_URL}/data/file_data" response = client.make_request( @@ -96,20 +95,20 @@ def __init__( :param kwargs: Additional file data (file_metadata, response, etc.) """ self.client = client - self.file_data = kwargs.get("file_data", {}) + self.__file_data = kwargs.get("file_data", {}) @property def file_id(self): - return self.file_data.get("file_id", "") + return self.__file_data.get("file_id", "") @property def project_id(self): - return self.file_data.get("project_id", "") + return self.__file_data.get("project_id", "") @property def dataset_id(self): - return self.file_data.get("dataset_id", "") + return self.__file_data.get("dataset_id", "") @property def metadata(self): - return self.file_data.get("file_metadata", {}) + return self.__file_data.get("file_metadata", {}) diff --git a/labellerr/core/files/video_file.py b/labellerr/core/files/video_file.py index 7e544d5..cb97956 100644 --- a/labellerr/core/files/video_file.py +++ b/labellerr/core/files/video_file.py @@ -63,7 +63,9 @@ def get_frames(self, frame_start: int = 0, frame_end: int | None = None): "uuid": unique_id, } - response = self.client.make_request(url, params, unique_id) + response = self.client.make_request( + "GET", url, extra_headers=None, request_id=unique_id, params=params + ) return response diff --git a/labellerr/core/projects/video_project.py b/labellerr/core/projects/video_project.py index cac694b..fd5bf15 100644 --- a/labellerr/core/projects/video_project.py +++ b/labellerr/core/projects/video_project.py @@ -2,18 +2,15 @@ from typing import List from .. import constants -from ..exceptions import LabellerrError from ..schemas import DatasetDataType, KeyFrame -from ..utils import validate_params from .base import LabellerrProject, LabellerrProjectMeta class VideoProject(LabellerrProject): """ - Class for handling video project operations and fetching multiple datasets. + Class for handling video project operations and fething multiple datasets. """ - @validate_params(file_id=str, keyframes=list) def add_or_update_keyframes( self, file_id: str, @@ -26,33 +23,26 @@ def add_or_update_keyframes( :param keyframes: List of KeyFrame objects to link :return: Response from the API """ - try: - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/actions/add_update_keyframes?client_id={self.client.client_id}&uuid={unique_id}" + unique_id = str(uuid.uuid4()) + url = f"{constants.BASE_URL}/actions/add_update_keyframes?client_id={self.client.client_id}&uuid={unique_id}" - body = { - "project_id": self.project_id, - "file_id": file_id, - "keyframes": [ - (kf.model_dump() if hasattr(kf, "model_dump") else kf) - for kf in keyframes - ], - } - - return self.client.make_request( - "POST", - url, - extra_headers={"content-type": "application/json"}, - request_id=unique_id, - json=body, - ) + body = { + "project_id": self.project_id, + "file_id": file_id, + "keyframes": [ + (kf.model_dump() if hasattr(kf, "model_dump") else kf) + for kf in keyframes + ], + } - except LabellerrError as e: - raise e - except Exception as e: - raise LabellerrError(f"Failed to link key frames: {str(e)}") + return self.client.make_request( + "POST", + url, + extra_headers={"content-type": "application/json"}, + request_id=unique_id, + json=body, + ) - @validate_params(file_id=str, keyframes=list) def delete_keyframes(self, file_id: str, keyframes: List[int]): """ Deletes key frames from a project. @@ -61,26 +51,20 @@ def delete_keyframes(self, file_id: str, keyframes: List[int]): :param keyframes: List of key frame numbers to delete :return: Response from the API """ - try: - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/actions/delete_keyframes?project_id={self.project_id}&uuid={unique_id}&client_id={self.client.client_id}" + unique_id = str(uuid.uuid4()) + url = f"{constants.BASE_URL}/actions/delete_keyframes?project_id={self.project_id}&uuid={unique_id}&client_id={self.client.client_id}" - return self.client.make_request( - "POST", - url, - extra_headers={"content-type": "application/json"}, - request_id=unique_id, - json={ - "project_id": self.project_id, - "file_id": file_id, - "keyframes": keyframes, - }, - ) - - except LabellerrError as e: - raise e - except Exception as e: - raise LabellerrError(f"Failed to delete key frames: {str(e)}") + return self.client.make_request( + "POST", + url, + extra_headers={"content-type": "application/json"}, + request_id=unique_id, + json={ + "project_id": self.project_id, + "file_id": file_id, + "keyframes": keyframes, + }, + ) LabellerrProjectMeta._register(DatasetDataType.video, VideoProject) diff --git a/labellerr/core/schemas/annotation_templates.py b/labellerr/core/schemas/annotation_templates.py index 7016015..b5c23f1 100644 --- a/labellerr/core/schemas/annotation_templates.py +++ b/labellerr/core/schemas/annotation_templates.py @@ -1,8 +1,8 @@ -from pydantic import BaseModel +from pydantic import BaseModel, Field from typing import List, Optional from enum import Enum from ..schemas import DatasetDataType - +import uuid class QuestionType(str, Enum): bounding_box = "BoundingBox" @@ -27,7 +27,7 @@ class AnnotationQuestion(BaseModel): question_number: int question: str - question_id: str + question_id: str = Field(default_factory=lambda: str(uuid.uuid4())) question_type: QuestionType required: bool options: Optional[List[Option]] = [] diff --git a/labellerr/core/utils.py b/labellerr/core/utils.py deleted file mode 100644 index 3d1727d..0000000 --- a/labellerr/core/utils.py +++ /dev/null @@ -1,140 +0,0 @@ -import logging -import time -from functools import wraps -from typing import Any, Callable, Optional, TypeVar, Union - -T = TypeVar("T") - - -def poll( - function: Callable[..., T], - condition: Callable[[T], bool], - interval: float = 2.0, - timeout: Optional[float] = None, - max_retries: Optional[int] = None, - args: tuple = (), - kwargs: dict = None, - on_success: Optional[Callable[[T], Any]] = None, - on_timeout: Optional[Callable[[int, Optional[T]], Any]] = None, - on_exception: Optional[Callable[[Exception], Any]] = None, -) -> Union[T, None]: - """ - Poll a function at specified intervals until a condition is met. - - Args: - function: The function to call - condition: Function that takes the return value of `function` and returns True when polling should stop - interval: Time in seconds between calls - timeout: Maximum time in seconds to poll before giving up - max_retries: Maximum number of retries before giving up - args: Positional arguments to pass to `function` - kwargs: Keyword arguments to pass to `function` - on_success: Callback function to call with the successful result - on_timeout: Callback function to call on timeout with the number of attempts and last result - on_exception: Callback function to call when an exception occurs in `function` - - Returns: - The last return value from `function` or None if timeout/max_retries was reached - - Examples: - ```python - # Poll until a job is complete - result = poll( - function=check_job_status, - condition=lambda status: status == "completed", - interval=5.0, - timeout=300, - args=(job_id,) - ) - - # Poll with a custom breaking condition - result = poll( - function=get_task_result, - condition=lambda r: r["status"] != "in_progress", - interval=2.0, - max_retries=10 - ) - ``` - """ - if kwargs is None: - kwargs = {} - - start_time = time.time() - attempts = 0 - last_result = None - - while True: - try: - attempts += 1 - last_result = function(*args, **kwargs) - - # Check if condition is satisfied - if condition(last_result): - if on_success: - on_success(last_result) - return last_result - - except Exception as e: - if on_exception: - on_exception(e) - logging.error(f"Exception in poll function: {str(e)}") - - # Check if we've reached timeout - if timeout is not None and time.time() - start_time > timeout: - if on_timeout: - on_timeout(attempts, last_result) - logging.warning( - f"Polling timed out after {timeout} seconds ({attempts} attempts)" - ) - return last_result - - # Check if we've reached max retries - if max_retries is not None and attempts >= max_retries: - if on_timeout: - on_timeout(attempts, last_result) - logging.warning(f"Polling reached max retries: {max_retries}") - return last_result - - # Wait before next attempt - time.sleep(interval) - - -def validate_params(**validations): - """ - Decorator to validate method parameters based on type specifications. - - Usage: - @validate_params(project_id=str, file_id=str, keyFrames=list) - def some_method(self, project_id, file_id, keyFrames): - ... - """ - - def decorator(func): - @wraps(func) - def wrapper(*args, **kwargs): - # Get function signature to map args to parameter names - import inspect - - sig = inspect.signature(func) - bound = sig.bind(*args, **kwargs) - bound.apply_defaults() - - # Validate each parameter - for param_name, expected_type in validations.items(): - if param_name in bound.arguments: - value = bound.arguments[param_name] - if not isinstance(value, expected_type): - from .exceptions import LabellerrError - - type_name = ( - " or ".join(t.__name__ for t in expected_type) - if isinstance(expected_type, tuple) - else expected_type.__name__ - ) - raise LabellerrError(f"{param_name} must be a {type_name}") - - return func(*args, **kwargs) - - return wrapper - - return decorator From 0b6078abda7f44a32d8adf3320aedc45f0f6f9b6 Mon Sep 17 00:00:00 2001 From: Yash Raj Suman Date: Wed, 12 Nov 2025 19:03:39 +0530 Subject: [PATCH 09/13] Feature/labimp 8041 (#27) * integration changes update * added test notebook * added integration test --- SDK_test.ipynb | 324 ++++++++++++++++++++++++ labellerr/core/constants.py | 2 +- labellerr/core/datasets/__init__.py | 4 +- labellerr/core/datasets/base.py | 13 +- labellerr/core/schemas/__init__.py | 10 + test/integration/test_whole_workflow.py | 109 ++++++++ 6 files changed, 451 insertions(+), 11 deletions(-) create mode 100644 SDK_test.ipynb create mode 100644 test/integration/test_whole_workflow.py diff --git a/SDK_test.ipynb b/SDK_test.ipynb new file mode 100644 index 0000000..8ef6d1b --- /dev/null +++ b/SDK_test.ipynb @@ -0,0 +1,324 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 6, + "id": "171b00fb", + "metadata": {}, + "outputs": [], + "source": [ + "from labellerr.client import LabellerrClient\n", + "from labellerr.core.datasets import create_dataset_from_local, LabellerrDataset\n", + "from labellerr.core.annotation_templates import create_template\n", + "from labellerr.core.projects import create_project\n", + "# from labellerr.core.schemas import * remove this code\n", + "from labellerr.core.schemas.annotation_templates import AnnotationQuestion, QuestionType, CreateTemplateParams\n", + "\n", + "import uuid\n" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "d0f76fc3", + "metadata": {}, + "outputs": [], + "source": [ + "from dotenv import dotenv_values\n", + "\n", + "config = dotenv_values(\".env\")\n", + "\n", + "API_KEY = config[\"API-KEY\"]\n", + "API_SECRET = config[\"API-SECRET\"]\n", + "CLIENT_ID = config[\"CLIENT_ID\"]" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "30bedb4f", + "metadata": {}, + "outputs": [], + "source": [ + "# Initialize client\n", + "client = LabellerrClient(api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID)\n" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "86549f09", + "metadata": {}, + "outputs": [], + "source": [ + "img_dataset_path = r\"D:\\Professional\\GitHub\\LABIMP-8041\\sample_img_dataset\"" + ] + }, + { + "cell_type": "code", + "execution_count": 15, + "id": "cd8b2606", + "metadata": {}, + "outputs": [], + "source": [ + "# 1. Create dataset\n", + "\n", + "dataset = create_dataset_from_local(\n", + " client=client,\n", + " dataset_config=DatasetConfig(dataset_name=\"My Dataset\", data_type=\"image\"),\n", + " folder_to_upload=img_dataset_path\n", + ")\n" + ] + }, + { + "cell_type": "code", + "execution_count": 16, + "id": "e25cd3c6", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'1c8b2a05-0321-44fd-91e3-2ea911382cf9'" + ] + }, + "execution_count": 16, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "dataset.dataset_id" + ] + }, + { + "cell_type": "raw", + "id": "f4f315e0", + "metadata": { + "vscode": { + "languageId": "raw" + } + }, + "source": [ + "'1c8b2a05-0321-44fd-91e3-2ea911382cf9'" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "35ffb634", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "DEBUG:urllib3.connectionpool:Starting new HTTPS connection (1): api.labellerr.com:443\n", + "DEBUG:urllib3.connectionpool:https://api.labellerr.com:443 \"GET /datasets/1c8b2a05-0321-44fd-91e3-2ea911382cf9?client_id=14836&uuid=dd77ee6e-a989-40dc-b32a-b467a02042cb HTTP/1.1\" 200 425\n", + "DEBUG:urllib3.connectionpool:https://api.labellerr.com:443 \"GET /datasets/1c8b2a05-0321-44fd-91e3-2ea911382cf9?client_id=14836&uuid=96f443fe-8f13-4971-907e-9ed0a2a37298 HTTP/1.1\" 200 425\n", + "INFO:root:Dataset 1c8b2a05-0321-44fd-91e3-2ea911382cf9 processing completed successfully!\n" + ] + }, + { + "data": { + "text/plain": [ + "{'es_multimodal_index': False,\n", + " 'metadata': {},\n", + " 'dataset_id': '1c8b2a05-0321-44fd-91e3-2ea911382cf9',\n", + " 'origin': 'https://pro.labellerr.com',\n", + " 'data_type': 'image',\n", + " 'name': 'My Dataset',\n", + " 'created_at': 1762930550910,\n", + " 'description': '',\n", + " 'created_by': 'a9555c.3c086748f3a839a246e7c3a60a',\n", + " 'client_id': '14836',\n", + " 'tags': [],\n", + " 'updated_at': 1762930676652,\n", + " 'progress': 'Processing 0/5 files',\n", + " 'files_count': 5,\n", + " 'status_code': 300,\n", + " 'es_index_status': 300}" + ] + }, + "execution_count": 5, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "import logging\n", + "logging.basicConfig(level=logging.DEBUG)\n", + "dataset = LabellerrDataset(\n", + " client=client, dataset_id=\"1c8b2a05-0321-44fd-91e3-2ea911382cf9\"\n", + ")\n", + "dataset.status()\n" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "afbef755", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "DEBUG:urllib3.connectionpool:Resetting dropped connection: api.labellerr.com\n", + "DEBUG:urllib3.util.retry:Incremented Retry for (url='/datasets/1c8b2a05-0321-44fd-91e3-2ea911382cf9?client_id=14836&uuid=b5ed59c7-1131-4c93-83b4-eafc39e6e27c'): Retry(total=2, connect=None, read=None, redirect=None, status=None)\n", + "WARNING:urllib3.connectionpool:Retrying (Retry(total=2, connect=None, read=None, redirect=None, status=None)) after connection broken by 'NameResolutionError(\": Failed to resolve 'api.labellerr.com' ([Errno 11001] getaddrinfo failed)\")': /datasets/1c8b2a05-0321-44fd-91e3-2ea911382cf9?client_id=14836&uuid=b5ed59c7-1131-4c93-83b4-eafc39e6e27c\n", + "DEBUG:urllib3.connectionpool:Starting new HTTPS connection (2): api.labellerr.com:443\n", + "DEBUG:urllib3.connectionpool:https://api.labellerr.com:443 \"GET /datasets/1c8b2a05-0321-44fd-91e3-2ea911382cf9?client_id=14836&uuid=b5ed59c7-1131-4c93-83b4-eafc39e6e27c HTTP/1.1\" 200 425\n", + "INFO:root:Dataset 1c8b2a05-0321-44fd-91e3-2ea911382cf9 processing completed successfully!\n" + ] + }, + { + "data": { + "text/plain": [ + "dict" + ] + }, + "execution_count": 12, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "type(dataset.status())" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "e1280726", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "DEBUG:urllib3.connectionpool:Resetting dropped connection: api.labellerr.com\n", + "DEBUG:urllib3.connectionpool:https://api.labellerr.com:443 \"POST /annotations/create_template?client_id=14836&data_type=image&uuid=426cd6b4-36c4-49de-8b3e-7ee6fc2c5a8b HTTP/1.1\" 200 67\n", + "DEBUG:urllib3.connectionpool:https://api.labellerr.com:443 \"GET /annotations/get_template?template_id=ee4b44ab-13a0-4f23-8b6a-f68bc235a46b&client_id=14836&uuid=4728db59-e7f7-41ed-8de1-d0525c8eba84 HTTP/1.1\" 200 422\n" + ] + } + ], + "source": [ + "from labellerr.core.schemas.annotation_templates import AnnotationQuestion, QuestionType, CreateTemplateParams\n", + "\n", + "# 2. Create annotation template\n", + "\n", + "template = create_template(\n", + " client=client,\n", + " params=CreateTemplateParams(\n", + " template_name=\"My Template\",\n", + " data_type=DatasetDataType.image,\n", + " questions=[\n", + " AnnotationQuestion(\n", + " question_number=1,\n", + " question=\"Object\",\n", + " question_id=str(uuid.uuid4()),\n", + " question_type=QuestionType.bounding_box,\n", + " required=True,\n", + " color=\"#FF0000\"\n", + " )\n", + " ]\n", + " )\n", + ")\n" + ] + }, + { + "cell_type": "code", + "execution_count": 11, + "id": "25af71c1", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "'ee4b44ab-13a0-4f23-8b6a-f68bc235a46b'" + ] + }, + "execution_count": 11, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "template.annotation_template_id" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "8d2c2a94", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "DEBUG:urllib3.connectionpool:https://api.labellerr.com:443 \"POST /projects/create?client_id=14836&uuid=57540ef2-204b-411a-9e66-79eda8d2da19 HTTP/1.1\" 200 96\n", + "DEBUG:urllib3.connectionpool:https://api.labellerr.com:443 \"GET /projects/project/roseline_neutral_perch_70115?client_id=14836&uuid=b6e284a0-13e7-494c-bcb1-98389f19a4b2 HTTP/1.1\" 200 729\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Project created: roseline_neutral_perch_70115\n" + ] + } + ], + "source": [ + "# 3. Create project\n", + "\n", + "project = create_project(\n", + " client=client,\n", + " params=CreateProjectParams(\n", + " project_name=\"My Project\",\n", + " data_type=DatasetDataType.image,\n", + " rotations=RotationConfig(\n", + " annotation_rotation_count=1,\n", + " review_rotation_count=1,\n", + " client_review_rotation_count=1\n", + " )\n", + " ),\n", + " datasets=[dataset],\n", + " annotation_template=template\n", + ")\n", + "\n", + "print(f\"✓ Project created: {project.project_id}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "daa2cebb", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.0" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/labellerr/core/constants.py b/labellerr/core/constants.py index 497e84d..5484189 100644 --- a/labellerr/core/constants.py +++ b/labellerr/core/constants.py @@ -1,4 +1,4 @@ -BASE_URL = "https://api-gateway-722091373895.us-central1.run.app" +BASE_URL = "https://api.labellerr.com" ALLOWED_ORIGINS = "https://pro.labellerr.com" diff --git a/labellerr/core/datasets/__init__.py b/labellerr/core/datasets/__init__.py index 46b1b18..d518c69 100644 --- a/labellerr/core/datasets/__init__.py +++ b/labellerr/core/datasets/__init__.py @@ -26,7 +26,7 @@ def create_dataset_from_connection( client: "LabellerrClient", dataset_config: schemas.DatasetConfig, - connection: LabellerrConnection, + connection: LabellerrConnection | str, path: str, ) -> LabellerrDataset: """ @@ -46,7 +46,7 @@ def create_dataset_from_connection( "dataset_name": dataset_config.dataset_name, "dataset_description": dataset_config.dataset_description, "data_type": dataset_config.data_type, - "connection_id": connection.connection_id, + "connection_id": connection.connection_id if isinstance(connection, LabellerrConnection) else connection, "path": path, "client_id": client.client_id, "es_multimodal_index": dataset_config.multimodal_indexing, diff --git a/labellerr/core/datasets/base.py b/labellerr/core/datasets/base.py index be01bfa..180e9a6 100644 --- a/labellerr/core/datasets/base.py +++ b/labellerr/core/datasets/base.py @@ -98,10 +98,7 @@ def data_type(self): return self.__dataset_data.get("data_type") def status( - self, - interval: float = 2.0, - timeout: Optional[float] = None, - max_retries: Optional[int] = None, + self ) -> Dict[str, Any]: """ Poll dataset status until completion or timeout. @@ -147,7 +144,7 @@ def get_dataset_status(): def is_completed(dataset_data): status_code = dataset_data.get("status_code", 500) # Consider dataset complete when status_code is 200 (success) or >= 400 (error/failed) - return status_code == 200 or status_code >= 400 + return status_code == 300 or status_code >= 400 def on_success(dataset_data): status_code = dataset_data.get("status_code", 500) @@ -166,9 +163,9 @@ def on_success(dataset_data): return poll( function=get_dataset_status, condition=is_completed, - interval=interval, - timeout=timeout, - max_retries=max_retries, + interval=2.0, + timeout=None, + max_retries=None, on_success=on_success, ) diff --git a/labellerr/core/schemas/__init__.py b/labellerr/core/schemas/__init__.py index 9fa4967..a3a9437 100644 --- a/labellerr/core/schemas/__init__.py +++ b/labellerr/core/schemas/__init__.py @@ -75,6 +75,16 @@ # Export schemas from labellerr.core.schemas.exports import CreateExportParams, ExportDestination + +# Export annotation templates +from labellerr.core.schemas.annotation_templates import ( + AnnotationQuestion, + Option, + QuestionType, +) + + + __all__ = [ # Base types "NonEmptyStr", diff --git a/test/integration/test_whole_workflow.py b/test/integration/test_whole_workflow.py new file mode 100644 index 0000000..16bdfa1 --- /dev/null +++ b/test/integration/test_whole_workflow.py @@ -0,0 +1,109 @@ +from labellerr.client import LabellerrClient +from labellerr.core.datasets import create_dataset_from_local +from labellerr.core.annotation_templates import create_template +from labellerr.core.projects import create_project +# from labellerr.core.schemas import * remove this code +from labellerr.core.schemas.annotation_templates import AnnotationQuestion, QuestionType, CreateTemplateParams +from labellerr.core.schemas.datasets import DatasetConfig +from labellerr.core.schemas import DatasetDataType +from labellerr.core.schemas.projects import CreateProjectParams, RotationConfig + +import uuid +import pytest +from dotenv import load_dotenv +import os + +load_dotenv() + +API_KEY = os.getenv("API_KEY") +API_SECRET = os.getenv("API_SECRET") +CLIENT_ID = os.getenv("CLIENT_ID") +IMG_DATASET_PATH = os.getenv("IMG_DATASET_PATH") + + +client = LabellerrClient(api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID) + +@pytest.fixture +def create_dataset_fixture(): + client = LabellerrClient(api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID) + + dataset = create_dataset_from_local( + client=client, + dataset_config=DatasetConfig(dataset_name="My Dataset", data_type="image"), + folder_to_upload=IMG_DATASET_PATH + ) + + return dataset + + +@pytest.fixture +def create_annotation_template_fixture(): + client = LabellerrClient(api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID) + + template = create_template( + client=client, + params=CreateTemplateParams( + template_name="My Template", + data_type=DatasetDataType.image, + questions=[ + AnnotationQuestion( + question_number=1, + question="Object", + question_id=str(uuid.uuid4()), + question_type=QuestionType.bounding_box, + required=True, + color="#FF0000" + ) + ] + ) + ) + + return template + +@pytest.fixture +def create_project_fixture(create_dataset_fixture, create_annotation_template_fixture): + client = LabellerrClient(api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID) + dataset = create_dataset_fixture + template = create_annotation_template_fixture + + project = create_project( + client=client, + params=CreateProjectParams( + project_name="My Project", + data_type=DatasetDataType.image, + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1 + ) + ), + datasets=[dataset], + annotation_template=template + ) + + return project + + +def test_create_dataset(create_dataset_fixture): + dataset = create_dataset_fixture + + assert dataset.dataset_id is not None + + result = dataset.status() + + assert result['status_code'] == 300 + assert result['files_count'] > 0 + + +def test_create_annotation_template(create_annotation_template_fixture): + template = create_annotation_template_fixture + + assert template.annotation_template_id is not None + assert isinstance(template.annotation_template_id, str) + + +def test_create_project(create_project_fixture): + project = create_project_fixture + + assert project.project_id is not None + assert isinstance(project.project_id, str) From 4b0833498c4411fbe65b19fb87d2528ebd9619e2 Mon Sep 17 00:00:00 2001 From: Yash Raj Suman Date: Thu, 13 Nov 2025 16:41:47 +0530 Subject: [PATCH 10/13] [LABIMP-8136 ] Add integration test on functions (#28) * integration changes update * added test notebook * added integration test * add integration tests for annotation templates, datasets, projects, and export functionality * refactor integration tests to use test-specific names for datasets and templates --- tests/integration/Export_project.py | 40 ------- .../test_create_annotation_template.py | 62 +++++++++++ tests/integration/test_create_dataset.py | 41 +++++++ tests/integration/test_create_project.py | 50 +++++++++ tests/integration/test_export_annotation.py | 54 ++++++++++ .../integration/test_whole_workflow.py | 100 ++++++++++-------- 6 files changed, 265 insertions(+), 82 deletions(-) delete mode 100644 tests/integration/Export_project.py create mode 100644 tests/integration/test_create_annotation_template.py create mode 100644 tests/integration/test_create_dataset.py create mode 100644 tests/integration/test_create_project.py create mode 100644 tests/integration/test_export_annotation.py rename {test => tests}/integration/test_whole_workflow.py (63%) diff --git a/tests/integration/Export_project.py b/tests/integration/Export_project.py deleted file mode 100644 index 6e0c1a3..0000000 --- a/tests/integration/Export_project.py +++ /dev/null @@ -1,40 +0,0 @@ -import os -import sys - -from labellerr import LabellerrError - -sys.path.append( - os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "SDKPython")) -) - -# Add the root directory to Python path -root_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) -sys.path.append(root_dir) - -from labellerr.client import LabellerrClient - - -# todo: ximi this don't use new struct -def export_project(api_key, api_secret, client_id, project_id): - """Exports a project using the Labellerr SDK.""" - - client = LabellerrClient(api_key, api_secret) - export_config = { - "export_name": "Weekly Export", - "export_description": "Export of all accepted annotations", - "export_format": "coco_json", - "statuses": [ - "review", - "r_assigned", - "client_review", - "cr_assigned", - "accepted", - ], - } - try: - result = client.create_local_export(project_id, client_id, export_config) - - export_id = result["response"]["report_id"] - print(f"Local export created successfully. Export ID: {export_id}") - except LabellerrError as e: - print(f"Local export creation failed: {str(e)}") diff --git a/tests/integration/test_create_annotation_template.py b/tests/integration/test_create_annotation_template.py new file mode 100644 index 0000000..09c233e --- /dev/null +++ b/tests/integration/test_create_annotation_template.py @@ -0,0 +1,62 @@ +import os +import uuid + +import pytest +from dotenv import load_dotenv + +from labellerr.client import LabellerrClient +from labellerr.core.annotation_templates import create_template +from labellerr.core.schemas import DatasetDataType +from labellerr.core.schemas.annotation_templates import ( + AnnotationQuestion, + CreateTemplateParams, + QuestionType, +) + +load_dotenv() + +API_KEY = os.getenv("API_KEY") +API_SECRET = os.getenv("API_SECRET") +CLIENT_ID = os.getenv("CLIENT_ID") + + +@pytest.fixture +def create_annotation_template_fixture(): + client = LabellerrClient( + api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID + ) + + template = create_template( + client=client, + params=CreateTemplateParams( + template_name="My Template", + data_type=DatasetDataType.image, + questions=[ + AnnotationQuestion( + question_number=1, + question="TEST QUESTION - Bounding Box", + question_id=str(uuid.uuid4()), + question_type=QuestionType.bounding_box, + required=True, + color="#FF0000", + ), + AnnotationQuestion( + question_number=2, + question="TEST QUESTION - Polygon", + question_id=str(uuid.uuid4()), + question_type=QuestionType.polygon, + required=True, + color="#FFC800", + ), + ], + ), + ) + + return template + + +def test_create_annotation_template(create_annotation_template_fixture): + template = create_annotation_template_fixture + + assert template.annotation_template_id is not None + assert isinstance(template.annotation_template_id, str) diff --git a/tests/integration/test_create_dataset.py b/tests/integration/test_create_dataset.py new file mode 100644 index 0000000..c373751 --- /dev/null +++ b/tests/integration/test_create_dataset.py @@ -0,0 +1,41 @@ +import os + +import pytest +from dotenv import load_dotenv + +from labellerr.client import LabellerrClient +from labellerr.core.datasets import create_dataset_from_local +from labellerr.core.schemas import DatasetConfig + +load_dotenv() + +API_KEY = os.getenv("API_KEY") +API_SECRET = os.getenv("API_SECRET") +CLIENT_ID = os.getenv("CLIENT_ID") +IMG_DATASET_PATH = os.getenv("IMG_DATASET_PATH") + + +@pytest.fixture +def create_dataset_fixture(): + client = LabellerrClient( + api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID + ) + + dataset = create_dataset_from_local( + client=client, + dataset_config=DatasetConfig(dataset_name="My Dataset", data_type="image"), + folder_to_upload=IMG_DATASET_PATH, + ) + + return dataset + + +def test_create_dataset(create_dataset_fixture): + dataset = create_dataset_fixture + + assert dataset.dataset_id is not None + + result = dataset.status() + + assert result["status_code"] == 300 + assert result["files_count"] > 0 diff --git a/tests/integration/test_create_project.py b/tests/integration/test_create_project.py new file mode 100644 index 0000000..ff3790b --- /dev/null +++ b/tests/integration/test_create_project.py @@ -0,0 +1,50 @@ +import os + +import pytest + +from labellerr.client import LabellerrClient +from labellerr.core.annotation_templates import LabellerrAnnotationTemplate +from labellerr.core.datasets import LabellerrDataset +from labellerr.core.projects import create_project +from labellerr.core.schemas import CreateProjectParams, DatasetDataType, RotationConfig + +API_KEY = os.getenv("API_KEY") +API_SECRET = os.getenv("API_SECRET") +CLIENT_ID = os.getenv("CLIENT_ID") +DATASET_ID = os.getenv("DATASET_ID") +TEMPLATE_ID = os.getenv("TEMPLATE_ID") + + +@pytest.fixture +def create_project_fixture(client): + + client = LabellerrClient( + api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID + ) + dataset = LabellerrDataset(client=client, dataset_id=DATASET_ID) + template = LabellerrAnnotationTemplate( + client=client, annotation_template_id=TEMPLATE_ID + ) + + project = create_project( + client=client, + params=CreateProjectParams( + project_name="My Project", + data_type=DatasetDataType.image, + rotations=RotationConfig( + annotation_rotation_count=1, + review_rotation_count=1, + client_review_rotation_count=1, + ), + ), + datasets=[dataset], + annotation_template=template, + ) + return project + + +def test_create_project(create_project_fixture): + project = create_project_fixture + + assert project.project_id is not None + assert isinstance(project.project_id, str) diff --git a/tests/integration/test_export_annotation.py b/tests/integration/test_export_annotation.py new file mode 100644 index 0000000..b04737c --- /dev/null +++ b/tests/integration/test_export_annotation.py @@ -0,0 +1,54 @@ +import os + +import pytest +from dotenv import load_dotenv + +from labellerr.client import LabellerrClient +from labellerr.core.projects import LabellerrProject + +load_dotenv() + +API_KEY = os.getenv("API_KEY") +API_SECRET = os.getenv("API_SECRET") +CLIENT_ID = os.getenv("CLIENT_ID") +PROJECT_ID = os.getenv("PROJECT_ID") + + +@pytest.fixture +def export_annotation_fixture(): + # Initialize the client with your API credentials + client = LabellerrClient( + api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID + ) + + project_id = PROJECT_ID + + export_config = { + "export_name": "Weekly Export", + "export_description": "Export of all accepted annotations", + "export_format": "coco_json", + "statuses": [ + "review", + "r_assigned", + "client_review", + "cr_assigned", + "accepted", + ], + } + + # Get project instance + project = LabellerrProject(client=client, project_id=project_id) + + # Create export + result = project.create_local_export(export_config) + export_id = result["response"]["report_id"] + # print(f"Local export created successfully. Export ID: {export_id}") + + return export_id + + +def test_export_annotation(export_annotation_fixture): + export_id = export_annotation_fixture + + assert export_id is not None + assert isinstance(export_id, str) diff --git a/test/integration/test_whole_workflow.py b/tests/integration/test_whole_workflow.py similarity index 63% rename from test/integration/test_whole_workflow.py rename to tests/integration/test_whole_workflow.py index 16bdfa1..fac44a8 100644 --- a/test/integration/test_whole_workflow.py +++ b/tests/integration/test_whole_workflow.py @@ -1,18 +1,24 @@ +import os +import uuid + +import pytest +from dotenv import load_dotenv + from labellerr.client import LabellerrClient -from labellerr.core.datasets import create_dataset_from_local from labellerr.core.annotation_templates import create_template +from labellerr.core.datasets import create_dataset_from_local from labellerr.core.projects import create_project +from labellerr.core.schemas import DatasetDataType + # from labellerr.core.schemas import * remove this code -from labellerr.core.schemas.annotation_templates import AnnotationQuestion, QuestionType, CreateTemplateParams +from labellerr.core.schemas.annotation_templates import ( + AnnotationQuestion, + CreateTemplateParams, + QuestionType, +) from labellerr.core.schemas.datasets import DatasetConfig -from labellerr.core.schemas import DatasetDataType from labellerr.core.schemas.projects import CreateProjectParams, RotationConfig -import uuid -import pytest -from dotenv import load_dotenv -import os - load_dotenv() API_KEY = os.getenv("API_KEY") @@ -21,89 +27,99 @@ IMG_DATASET_PATH = os.getenv("IMG_DATASET_PATH") -client = LabellerrClient(api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID) +CLIENT = LabellerrClient(api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID) + @pytest.fixture -def create_dataset_fixture(): - client = LabellerrClient(api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID) - +def create_dataset_fixture(client=CLIENT): + # client = LabellerrClient( + # api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID + # ) + dataset = create_dataset_from_local( - client=client, - dataset_config=DatasetConfig(dataset_name="My Dataset", data_type="image"), - folder_to_upload=IMG_DATASET_PATH + client=client, + dataset_config=DatasetConfig(dataset_name="TEST DATASET", data_type="image"), + folder_to_upload=IMG_DATASET_PATH, ) - + return dataset @pytest.fixture -def create_annotation_template_fixture(): - client = LabellerrClient(api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID) - +def create_annotation_template_fixture(client=CLIENT): + # client = LabellerrClient( + # api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID + # ) + template = create_template( client=client, params=CreateTemplateParams( - template_name="My Template", + template_name="TEST TEMPLATE", data_type=DatasetDataType.image, questions=[ AnnotationQuestion( question_number=1, - question="Object", + question="TEST QUESTION - Bounding Box", question_id=str(uuid.uuid4()), question_type=QuestionType.bounding_box, required=True, - color="#FF0000" + color="#FF0000", ) - ] - ) + ], + ), ) - + return template + @pytest.fixture -def create_project_fixture(create_dataset_fixture, create_annotation_template_fixture): - client = LabellerrClient(api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID) +def create_project_fixture( + create_dataset_fixture, create_annotation_template_fixture, client=CLIENT +): + # client = LabellerrClient( + # api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID + # ) dataset = create_dataset_fixture template = create_annotation_template_fixture - + project = create_project( client=client, params=CreateProjectParams( - project_name="My Project", + project_name="TEST PROJECT", data_type=DatasetDataType.image, rotations=RotationConfig( annotation_rotation_count=1, review_rotation_count=1, - client_review_rotation_count=1 - ) + client_review_rotation_count=1, + ), ), datasets=[dataset], - annotation_template=template + annotation_template=template, ) - + return project def test_create_dataset(create_dataset_fixture): dataset = create_dataset_fixture - + assert dataset.dataset_id is not None - + result = dataset.status() - - assert result['status_code'] == 300 - assert result['files_count'] > 0 - - + + assert result["status_code"] == 300 + assert result["files_count"] > 0 + + def test_create_annotation_template(create_annotation_template_fixture): template = create_annotation_template_fixture - + assert template.annotation_template_id is not None assert isinstance(template.annotation_template_id, str) def test_create_project(create_project_fixture): project = create_project_fixture - + assert project.project_id is not None assert isinstance(project.project_id, str) From 7f28ec9c7eded88292001cf4cd5ed1926c73d195 Mon Sep 17 00:00:00 2001 From: Ximi Hoque Date: Thu, 13 Nov 2025 17:26:27 +0530 Subject: [PATCH 11/13] Updates to sync datasets --- drivers/connectors.py | 73 ------------------- drivers/datasets.py | 64 ---------------- drivers/projects.py | 65 ----------------- labellerr/core/connectors/connections.py | 4 +- labellerr/core/datasets/__init__.py | 6 +- labellerr/core/datasets/base.py | 36 +++++---- labellerr/core/projects/base.py | 6 +- labellerr/core/schemas/__init__.py | 5 +- .../core/schemas/annotation_templates.py | 1 + 9 files changed, 37 insertions(+), 223 deletions(-) delete mode 100644 drivers/connectors.py delete mode 100644 drivers/datasets.py delete mode 100644 drivers/projects.py diff --git a/drivers/connectors.py b/drivers/connectors.py deleted file mode 100644 index f5d51d5..0000000 --- a/drivers/connectors.py +++ /dev/null @@ -1,73 +0,0 @@ -import logging -import os -from dotenv import load_dotenv -from labellerr.client import LabellerrClient - -# from labellerr.core.connectors import ( -# LabellerrConnection, -# list_connections, -# delete_connection, -# create_connection, -# AWSConnectionParams, -# LabellerrGCSConnection, -# ) -# from labellerr.core.schemas import ( -# ConnectionType, -# ConnectorType, -# DatasetDataType, -# GCSConnectionTestParams, -# GCSConnectionParams, -# ) - -# Set logging level to DEBUG -logging.basicConfig(level=logging.DEBUG) - -load_dotenv() - -API_KEY = os.getenv("API_KEY") -API_SECRET = os.getenv("API_SECRET") -CLIENT_ID = os.getenv("CLIENT_ID") - -if not all([API_KEY, API_SECRET, CLIENT_ID]): - raise ValueError( - "API_KEY, API_SECRET, and CLIENT_ID must be set in environment variables" - ) - -# Initialize client -client = LabellerrClient( - api_key=API_KEY, - api_secret=API_SECRET, - client_id=CLIENT_ID, -) - -# response = LabellerrGCSConnection.test_connection(client=client, params=GCSConnectionTestParams( -# svc_account_json='labellerr-dev.json', -# path="gs://aws-labellerr-public-datasets/coco2017", -# connection_type=ConnectionType._IMPORT, -# data_type=DatasetDataType.video, -# )) -# print(response) - -# response = create_connection(client=client, connector_type=ConnectorType._GCS, params=GCSConnectionParams( -# svc_account_json='labellerr-dev.json', -# path="gs://aws-labellerr-public-datasets/coco2017", -# connection_type=ConnectionType._IMPORT, -# data_type=DatasetDataType.image, name="GCS Import Test", description="GCS Import Test")) -# print(response) -# connection = LabellerrConnection(client=client, connection_id='8c3dc4b4-e701-4d22-add3-28abc33e13ef') -# response = connection.test(path="gs://aws-labellerr-public-datasets/coco2017", connection_type=ConnectionType._IMPORT, data_type=DatasetDataType.image) -# print(response) -# connection = create_connection(client=client, connector_type=ConnectorType._S3, params=AWSConnectionParams( -# aws_access_key=os.getenv("AWS_KEY"), -# aws_secrets_key=os.getenv("AWS_SECRET"), -# s3_path="s3://amazon-s3-sync-test/labellerr-processed/videos/", # this path is not part of the connection but needed to test the connection on the desired path. -# # This can be dynamically changed while using the connection for creating datasets. -# connection_type=ConnectionType._IMPORT, -# name="Amazon S3 Import Test", -# description="Amazon S3 Import Test", -# )) -# print(f"Connection created: {connection.connection_id}") -# connection = LabellerrConnection(client=client, connection_id='2a30c044-57f7-42c7-8290-bab3bbac0ebc') -# print('connection type', connection.connection_type) - -# print (connection.test(s3_path="s3://amazon-s3-sync-test/labellerr-processed/videos/datasets/", connection_type=ConnectionType._IMPORT)) diff --git a/drivers/datasets.py b/drivers/datasets.py deleted file mode 100644 index f57356e..0000000 --- a/drivers/datasets.py +++ /dev/null @@ -1,64 +0,0 @@ -# from labellerr.core.datasets import ( -# create_dataset_from_local, -# create_dataset_from_connection, -# ) -import logging -import os -from dotenv import load_dotenv -from labellerr.client import LabellerrClient -from labellerr.core.schemas import DatasetConfig -from labellerr.core.datasets import LabellerrDataset - -# from labellerr.core.connectors import ( -# LabellerrConnection, -# list_connections, -# delete_connection, -# create_connection, -# AWSConnectionParams, -# LabellerrGCSConnection, -# ) -# from labellerr.core.schemas import ( -# ConnectionType, -# ConnectorType, -# DatasetDataType, -# GCSConnectionTestParams, -# GCSConnectionParams, -# ) - -# Set logging level to DEBUG -logging.basicConfig(level=logging.DEBUG) - -load_dotenv() - -API_KEY = os.getenv("API_KEY") -API_SECRET = os.getenv("API_SECRET") -CLIENT_ID = os.getenv("CLIENT_ID") - -if not all([API_KEY, API_SECRET, CLIENT_ID]): - raise ValueError( - "API_KEY, API_SECRET, and CLIENT_ID must be set in environment variables" - ) - -# Initialize client -client = LabellerrClient( - api_key=API_KEY, - api_secret=API_SECRET, - client_id=CLIENT_ID, -) -dataset_config = DatasetConfig( - dataset_name="test_dataset_from_local SDK", - dataset_description="test dataset description", - data_type="image", -) -# dataset = create_dataset_from_local( -# client=client, -# dataset_config=dataset_config, -# folder_to_upload='images_single', -# ) -dataset = LabellerrDataset( - client=client, dataset_id="455e3d45-55f9-436d-98c2-07a514b7894e" -) -print(dataset.files_count) - -for file in dataset.fetch_files(): - print(file.file_id) diff --git a/drivers/projects.py b/drivers/projects.py deleted file mode 100644 index 1afb013..0000000 --- a/drivers/projects.py +++ /dev/null @@ -1,65 +0,0 @@ -# from labellerr.core.projects import LabellerrProject -# from labellerr.core.annotation_templates import LabellerrAnnotationTemplate -import logging -import os -from dotenv import load_dotenv -from labellerr.client import LabellerrClient - -# from labellerr.core.schemas.annotation_templates import CreateTemplateParams, AnnotationQuestion, QuestionType -# from labellerr.core.annotation_templates import create_template -# from labellerr.core.schemas.projects import CreateProjectParams, RotationConfig -# from labellerr.core.projects import create_project, list_projects -# from labellerr.core.datasets import LabellerrDataset - - -# Set logging level to DEBUG -logging.basicConfig(level=logging.DEBUG) - -load_dotenv() - -API_KEY = os.getenv("API_KEY") -API_SECRET = os.getenv("API_SECRET") -CLIENT_ID = os.getenv("CLIENT_ID") - -if not all([API_KEY, API_SECRET, CLIENT_ID]): - raise ValueError( - "API_KEY, API_SECRET, and CLIENT_ID must be set in environment variables" - ) - -# Initialize client -client = LabellerrClient( - api_key=API_KEY, - api_secret=API_SECRET, - client_id=CLIENT_ID, -) - -# project = LabellerrProject(client=client, project_id="rafaela_youngest_pike_23125") -# res = project.upload_preannotation(annotation_format="coco_json", annotation_file="/Users/Ximi-Hoque/Downloads/export_to_annotate_05_15.json").result() -# print(res) - -# annotation_template = LabellerrAnnotationTemplate(client=client, annotation_template_id="00016829-9051-46b1-96c6-3ec6763c342a") -# print(annotation_template.annotation_template_data) - -# res = create_template(client, CreateTemplateParams(template_name="test_template_1", data_type="image", -# questions=[ -# AnnotationQuestion(question_number=1, -# question="test_question", -# question_id="test_question_id", -# question_type=QuestionType.bounding_box, -# required=True, -# color="#FF4500", -# )])) -# print(res) - -# project = create_project(client, -# CreateProjectParams( -# project_name="test_project_via_sdk", -# data_type="image", -# rotations=RotationConfig(annotation_rotation_count=1, review_rotation_count=1, client_review_rotation_count=1), -# use_ai=False), -# datasets=[LabellerrDataset(client=client, dataset_id="ca298293-7f5e-4bdd-801f-8863a5ba458b")], -# annotation_template=res -# ) -# print (project.project_data) -# project = LabellerrProject(client=client, project_id="dinnie_confidential_lynx_20766") -# print(project.project_data) diff --git a/labellerr/core/connectors/connections.py b/labellerr/core/connectors/connections.py index 0c56f27..2b9bc1b 100644 --- a/labellerr/core/connectors/connections.py +++ b/labellerr/core/connectors/connections.py @@ -62,7 +62,7 @@ class LabellerrConnection(metaclass=LabellerrConnectionMeta): def __init__(self, client: "LabellerrClient", connection_id: str, **kwargs): self.client = client - self._connection_id_input = connection_id + self.__connection_id_input = connection_id self.__connection_data = kwargs["connection_data"] @property @@ -75,7 +75,7 @@ def description(self): @property def connection_id(self): - return self.__connection_data.get("connection_id") + return self.__connection_id_input @property def connection_type(self): diff --git a/labellerr/core/datasets/__init__.py b/labellerr/core/datasets/__init__.py index d518c69..f02bfdf 100644 --- a/labellerr/core/datasets/__init__.py +++ b/labellerr/core/datasets/__init__.py @@ -46,7 +46,11 @@ def create_dataset_from_connection( "dataset_name": dataset_config.dataset_name, "dataset_description": dataset_config.dataset_description, "data_type": dataset_config.data_type, - "connection_id": connection.connection_id if isinstance(connection, LabellerrConnection) else connection, + "connection_id": ( + connection.connection_id + if isinstance(connection, LabellerrConnection) + else connection + ), "path": path, "client_id": client.client_id, "es_multimodal_index": dataset_config.multimodal_indexing, diff --git a/labellerr/core/datasets/base.py b/labellerr/core/datasets/base.py index 180e9a6..b41355f 100644 --- a/labellerr/core/datasets/base.py +++ b/labellerr/core/datasets/base.py @@ -4,13 +4,17 @@ import logging import uuid from abc import ABCMeta -from typing import Dict, Optional, Any, List +from typing import Dict, Any, List, TYPE_CHECKING from .. import constants from ..exceptions import InvalidDatasetError, LabellerrError from ..client import LabellerrClient from ..files import LabellerrFile +from ..connectors import LabellerrConnection + +if TYPE_CHECKING: + from ..projects import LabellerrProject class LabellerrDatasetMeta(ABCMeta): @@ -66,9 +70,13 @@ class LabellerrDataset(metaclass=LabellerrDatasetMeta): def __init__(self, client: "LabellerrClient", dataset_id: str, **kwargs): self.client = client - self.dataset_id = dataset_id + self.__dataset_id_input = dataset_id self.__dataset_data = kwargs["dataset_data"] + @property + def dataset_id(self): + return self.__dataset_id_input + @property def name(self): return self.__dataset_data.get("name") @@ -97,9 +105,7 @@ def status_code(self): def data_type(self): return self.__dataset_data.get("data_type") - def status( - self - ) -> Dict[str, Any]: + def status(self) -> Dict[str, Any]: """ Poll dataset status until completion or timeout. @@ -235,20 +241,18 @@ def fetch_files(self, page_size: int = 1000) -> List[LabellerrFile]: def sync_with_connection( self, - project_id, - path, - data_type, - email_id, - connection_id, + project: "LabellerrProject", + path: str, + data_type: str, + connection: LabellerrConnection, ): """ Syncs datasets with the backend. - :param project_id: The ID of the project + :param project: The project instance :param path: The path to sync :param data_type: Type of data (image, video, audio, document, text) - :param email_id: Email ID of the user - :param connection_id: The connection ID + :param connection: The connection instance :return: Dictionary containing sync status :raises LabellerrError: If the sync fails """ @@ -259,12 +263,12 @@ def sync_with_connection( payload = json.dumps( { "client_id": self.client.client_id, - "project_id": project_id, + "project_id": project.project_id, "dataset_id": self.dataset_id, "path": path, "data_type": data_type, - "email_id": email_id, - "connection_id": connection_id, + "email_id": self.client.api_key, + "connection_id": connection.connection_id, } ) diff --git a/labellerr/core/projects/base.py b/labellerr/core/projects/base.py index 14e05de..3ad5210 100644 --- a/labellerr/core/projects/base.py +++ b/labellerr/core/projects/base.py @@ -71,9 +71,13 @@ class LabellerrProject(metaclass=LabellerrProjectMeta): def __init__(self, client: "LabellerrClient", project_id: str, **kwargs): self.client = client - self.project_id = project_id + self.__project_id_input = project_id self.__project_data = kwargs["project_data"] + @property + def project_id(self): + return self.__project_id_input + @property def status_code(self): return self.__project_data.get("status_code", 501) # if not found, return 501 diff --git a/labellerr/core/schemas/__init__.py b/labellerr/core/schemas/__init__.py index a3a9437..5579571 100644 --- a/labellerr/core/schemas/__init__.py +++ b/labellerr/core/schemas/__init__.py @@ -84,7 +84,6 @@ ) - __all__ = [ # Base types "NonEmptyStr", @@ -133,4 +132,8 @@ # Export schemas "CreateExportParams", "ExportDestination", + # Annotation templates schemas + "AnnotationQuestion", + "Option", + "QuestionType", ] diff --git a/labellerr/core/schemas/annotation_templates.py b/labellerr/core/schemas/annotation_templates.py index b5c23f1..4885737 100644 --- a/labellerr/core/schemas/annotation_templates.py +++ b/labellerr/core/schemas/annotation_templates.py @@ -4,6 +4,7 @@ from ..schemas import DatasetDataType import uuid + class QuestionType(str, Enum): bounding_box = "BoundingBox" polygon = "polygon" From fc507963043b47654299e22e60687f58e0ba729f Mon Sep 17 00:00:00 2001 From: Ximi Hoque Date: Thu, 13 Nov 2025 17:26:51 +0530 Subject: [PATCH 12/13] Updated tests --- tests/integration/test_whole_workflow.py | 57 +++++++++++------------- 1 file changed, 27 insertions(+), 30 deletions(-) diff --git a/tests/integration/test_whole_workflow.py b/tests/integration/test_whole_workflow.py index fac44a8..c94c2f0 100644 --- a/tests/integration/test_whole_workflow.py +++ b/tests/integration/test_whole_workflow.py @@ -1,65 +1,64 @@ -import os -import uuid - -import pytest -from dotenv import load_dotenv - from labellerr.client import LabellerrClient -from labellerr.core.annotation_templates import create_template from labellerr.core.datasets import create_dataset_from_local +from labellerr.core.annotation_templates import create_template from labellerr.core.projects import create_project -from labellerr.core.schemas import DatasetDataType # from labellerr.core.schemas import * remove this code from labellerr.core.schemas.annotation_templates import ( AnnotationQuestion, - CreateTemplateParams, QuestionType, + CreateTemplateParams, ) from labellerr.core.schemas.datasets import DatasetConfig +from labellerr.core.schemas import DatasetDataType from labellerr.core.schemas.projects import CreateProjectParams, RotationConfig +import uuid +import pytest +from dotenv import load_dotenv +import os + load_dotenv() API_KEY = os.getenv("API_KEY") API_SECRET = os.getenv("API_SECRET") CLIENT_ID = os.getenv("CLIENT_ID") -IMG_DATASET_PATH = os.getenv("IMG_DATASET_PATH") +IMAGES_DATASET_UPLOAD_PATH = os.getenv("IMAGES_DATASET_UPLOAD_PATH") -CLIENT = LabellerrClient(api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID) +client = LabellerrClient(api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID) @pytest.fixture -def create_dataset_fixture(client=CLIENT): - # client = LabellerrClient( - # api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID - # ) +def create_dataset_fixture(): + client = LabellerrClient( + api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID + ) dataset = create_dataset_from_local( client=client, - dataset_config=DatasetConfig(dataset_name="TEST DATASET", data_type="image"), - folder_to_upload=IMG_DATASET_PATH, + dataset_config=DatasetConfig(dataset_name="My Dataset", data_type="image"), + folder_to_upload=IMAGES_DATASET_UPLOAD_PATH, ) return dataset @pytest.fixture -def create_annotation_template_fixture(client=CLIENT): - # client = LabellerrClient( - # api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID - # ) +def create_annotation_template_fixture(): + client = LabellerrClient( + api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID + ) template = create_template( client=client, params=CreateTemplateParams( - template_name="TEST TEMPLATE", + template_name="My Template", data_type=DatasetDataType.image, questions=[ AnnotationQuestion( question_number=1, - question="TEST QUESTION - Bounding Box", + question="Object", question_id=str(uuid.uuid4()), question_type=QuestionType.bounding_box, required=True, @@ -73,19 +72,17 @@ def create_annotation_template_fixture(client=CLIENT): @pytest.fixture -def create_project_fixture( - create_dataset_fixture, create_annotation_template_fixture, client=CLIENT -): - # client = LabellerrClient( - # api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID - # ) +def create_project_fixture(create_dataset_fixture, create_annotation_template_fixture): + client = LabellerrClient( + api_key=API_KEY, api_secret=API_SECRET, client_id=CLIENT_ID + ) dataset = create_dataset_fixture template = create_annotation_template_fixture project = create_project( client=client, params=CreateProjectParams( - project_name="TEST PROJECT", + project_name="My Project", data_type=DatasetDataType.image, rotations=RotationConfig( annotation_rotation_count=1, From 910cd5f54e9c3c11dac6deef82ae59dccb85a40d Mon Sep 17 00:00:00 2001 From: Ximi Hoque Date: Thu, 13 Nov 2025 17:38:04 +0530 Subject: [PATCH 13/13] Updated tests --- labellerr/core/projects/video_project.py | 91 +++++++++++++++--------- tests/unit/test_client.py | 4 +- tests/unit/test_create_dataset_path.py | 18 ++--- tests/unit/test_keyframes.py | 22 +++--- 4 files changed, 77 insertions(+), 58 deletions(-) diff --git a/labellerr/core/projects/video_project.py b/labellerr/core/projects/video_project.py index fd5bf15..a70d21c 100644 --- a/labellerr/core/projects/video_project.py +++ b/labellerr/core/projects/video_project.py @@ -2,6 +2,7 @@ from typing import List from .. import constants +from ..exceptions import LabellerrError from ..schemas import DatasetDataType, KeyFrame from .base import LabellerrProject, LabellerrProjectMeta @@ -23,25 +24,37 @@ def add_or_update_keyframes( :param keyframes: List of KeyFrame objects to link :return: Response from the API """ - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/actions/add_update_keyframes?client_id={self.client.client_id}&uuid={unique_id}" - - body = { - "project_id": self.project_id, - "file_id": file_id, - "keyframes": [ - (kf.model_dump() if hasattr(kf, "model_dump") else kf) - for kf in keyframes - ], - } - - return self.client.make_request( - "POST", - url, - extra_headers={"content-type": "application/json"}, - request_id=unique_id, - json=body, - ) + # Parameter validation + if not isinstance(file_id, str): + raise LabellerrError("file_id must be a str") + + if not isinstance(keyframes, list): + raise LabellerrError("keyframes must be a list") + + try: + unique_id = str(uuid.uuid4()) + url = f"{constants.BASE_URL}/actions/add_update_keyframes?client_id={self.client.client_id}&uuid={unique_id}" + + body = { + "project_id": self.project_id, + "file_id": file_id, + "keyframes": [ + (kf.model_dump() if hasattr(kf, "model_dump") else kf) + for kf in keyframes + ], + } + + return self.client.make_request( + "POST", + url, + extra_headers={"content-type": "application/json"}, + request_id=unique_id, + json=body, + ) + except LabellerrError: + raise + except Exception as e: + raise LabellerrError(f"Failed to link key frames: {str(e)}") def delete_keyframes(self, file_id: str, keyframes: List[int]): """ @@ -51,20 +64,32 @@ def delete_keyframes(self, file_id: str, keyframes: List[int]): :param keyframes: List of key frame numbers to delete :return: Response from the API """ - unique_id = str(uuid.uuid4()) - url = f"{constants.BASE_URL}/actions/delete_keyframes?project_id={self.project_id}&uuid={unique_id}&client_id={self.client.client_id}" - - return self.client.make_request( - "POST", - url, - extra_headers={"content-type": "application/json"}, - request_id=unique_id, - json={ - "project_id": self.project_id, - "file_id": file_id, - "keyframes": keyframes, - }, - ) + # Parameter validation + if not isinstance(file_id, str): + raise LabellerrError("file_id must be a str") + + if not isinstance(keyframes, list): + raise LabellerrError("keyframes must be a list") + + try: + unique_id = str(uuid.uuid4()) + url = f"{constants.BASE_URL}/actions/delete_keyframes?project_id={self.project_id}&uuid={unique_id}&client_id={self.client.client_id}" + + return self.client.make_request( + "POST", + url, + extra_headers={"content-type": "application/json"}, + request_id=unique_id, + json={ + "project_id": self.project_id, + "file_id": file_id, + "keyframes": keyframes, + }, + ) + except LabellerrError: + raise + except Exception as e: + raise LabellerrError(f"Failed to delete key frames: {str(e)}") LabellerrProjectMeta._register(DatasetDataType.video, VideoProject) diff --git a/tests/unit/test_client.py b/tests/unit/test_client.py index 96df217..91923eb 100644 --- a/tests/unit/test_client.py +++ b/tests/unit/test_client.py @@ -30,8 +30,8 @@ def project(client): # Use __new__ to create instance without calling __init__ through metaclass proj = ImageProject.__new__(ImageProject) proj.client = client - proj.project_id = "test_project_id" - proj.__project_data = project_data + proj._LabellerrProject__project_id_input = "test_project_id" + proj._LabellerrProject__project_data = project_data return proj diff --git a/tests/unit/test_create_dataset_path.py b/tests/unit/test_create_dataset_path.py index eee98a7..50c5f61 100644 --- a/tests/unit/test_create_dataset_path.py +++ b/tests/unit/test_create_dataset_path.py @@ -41,16 +41,13 @@ def test_create_dataset_from_connection_success(self, client): "labellerr.core.datasets.base.LabellerrDataset.get_dataset", return_value={"dataset_id": "test-dataset-id", "data_type": "image"}, ): - # Create a mock connection object - from unittest.mock import Mock - - mock_connection = Mock() - mock_connection.connection_id = "test-connection-id" + # Use string connection_id instead of Mock object + connection_id = "test-connection-id" dataset = create_dataset_from_connection( client=client, dataset_config=dataset_config, - connection=mock_connection, + connection=connection_id, path="s3://test-bucket/path/to/data", ) @@ -177,16 +174,13 @@ def test_create_dataset_from_connection_with_different_paths(self, client): "labellerr.core.datasets.base.LabellerrDataset.get_dataset", return_value={"dataset_id": "test-dataset-id", "data_type": "image"}, ): - # Create a mock connection object - from unittest.mock import Mock - - mock_connection = Mock() - mock_connection.connection_id = "test-gcp-connection-id" + # Use string connection_id instead of Mock object + connection_id = "test-gcp-connection-id" dataset = create_dataset_from_connection( client=client, dataset_config=dataset_config, - connection=mock_connection, + connection=connection_id, path="gs://test-bucket/path/to/data", ) diff --git a/tests/unit/test_keyframes.py b/tests/unit/test_keyframes.py index 9650254..d3a62b5 100644 --- a/tests/unit/test_keyframes.py +++ b/tests/unit/test_keyframes.py @@ -5,7 +5,7 @@ validation decorators, and keyframe-related client methods. """ -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from pydantic import ValidationError @@ -200,8 +200,8 @@ def mock_video_project(mock_client): # Create instance bypassing metaclass project = VideoProject.__new__(VideoProject) project.client = mock_client - project.project_id = "test_project_id" - project.__project_data = { + project._LabellerrProject__project_id_input = "test_project_id" + project._LabellerrProject__project_data = { "project_id": "test_project_id", "data_type": "video", "attached_datasets": [], @@ -298,12 +298,11 @@ def test_add_or_update_keyframes_invalid_parameters( with pytest.raises(LabellerrError, match=expected_error): mock_video_project.add_or_update_keyframes(file_id, keyframes) - @patch("labellerr.core.client.LabellerrClient.make_request") - def test_add_or_update_keyframes_api_error( - self, mock_make_request, mock_video_project - ): + def test_add_or_update_keyframes_api_error(self, mock_video_project): """Test add_or_update_keyframes when API call fails""" - mock_make_request.side_effect = Exception("API Error") + mock_video_project.client.make_request = MagicMock( + side_effect=Exception("API Error") + ) keyframes = [KeyFrame(frame_number=0)] with pytest.raises( @@ -391,10 +390,11 @@ def test_delete_keyframes_invalid_parameters( with pytest.raises(LabellerrError, match=expected_error): mock_video_project.delete_keyframes(file_id, keyframes) - @patch("labellerr.core.client.LabellerrClient.make_request") - def test_delete_keyframes_api_error(self, mock_make_request, mock_video_project): + def test_delete_keyframes_api_error(self, mock_video_project): """Test delete_keyframes when API call fails""" - mock_make_request.side_effect = Exception("API Error") + mock_video_project.client.make_request = MagicMock( + side_effect=Exception("API Error") + ) with pytest.raises( LabellerrError, match="Failed to delete key frames: API Error"