Changes

Author: midhun-pm

libs/labelbox/src/labelbox/alignerr/alignerr_project.py

@@ -5,6 +5,7 @@
 
 from labelbox.alignerr.schema.project_rate import ProjectRateV2
 from labelbox.alignerr.schema.project_domain import ProjectDomain
+from labelbox.alignerr.schema.enchanced_resource_tags import EnhancedResourceTag, ResourceTagType
 from labelbox.pagination import PaginatedCollection
 
 logger = logging.getLogger(__name__)
@@ -43,11 +44,6 @@ def project(self, project: "Project"):
         self._project = project
 
     def domains(self) -> PaginatedCollection:
-        """Get all domains associated with this project.
-
-        Returns:
-            PaginatedCollection of ProjectDomain instances
-        """
         return ProjectDomain.get_by_project_id(
             client=self.client, project_id=self.project.uid
         )
@@ -59,7 +55,7 @@ def add_domain(self, project_domain: ProjectDomain):
             domain_ids=[project_domain.uid],
         )
 
-    def get_project_rate(self) -> Optional["ProjectRateV2"]:
+    def get_project_rates(self) -> list["ProjectRateV2"]:
         return ProjectRateV2.get_by_project_id(
             client=self.client, project_id=self.project.uid
         )
@@ -71,6 +67,71 @@ def set_project_rate(self, project_rate_input):
             project_rate_input=project_rate_input,
         )
 
+    def set_tags(self, tag_names: list[str], tag_type: ResourceTagType):
+        # Convert tag names to tag IDs
+        tag_ids = []
+        for tag_name in tag_names:
+            # Search for the tag by text to get its ID
+            found_tags = EnhancedResourceTag.search_by_text(self.client, search_text=tag_name, tag_type=tag_type)
+            if found_tags:
+                tag_ids.append(found_tags[0].id)
+        
+        # Use the existing project resource tag functionality with IDs
+        self.project.update_project_resource_tags(tag_ids)
+        return self
+
+    def get_tags(self) -> list[EnhancedResourceTag]:
+        """Get enhanced resource tags associated with this project.
+        
+        Returns:
+            List of EnhancedResourceTag instances
+        """
+        # Get project resource tags and convert to EnhancedResourceTag instances
+        project_resource_tags = self.project.get_resource_tags()
+        enhanced_tags = []
+        for tag in project_resource_tags:
+            # Search for the corresponding EnhancedResourceTag by text (try different types)
+            found_tags = []
+            for tag_type in [ResourceTagType.Default, ResourceTagType.Billing]:
+                found_tags = EnhancedResourceTag.search_by_text(self.client, search_text=tag.text, tag_type=tag_type)
+                if found_tags:
+                    break
+            if found_tags:
+                enhanced_tags.extend(found_tags)
+        return enhanced_tags
+
+    def add_tag(self, tag: EnhancedResourceTag):
+        """Add a single enhanced resource tag to the project.
+        
+        Args:
+            tag: EnhancedResourceTag instance to add
+            
+        Returns:
+            Self for method chaining
+        """
+        current_tags = self.get_tags()
+        current_tag_names = [t.text for t in current_tags]
+        
+        if tag.text not in current_tag_names:
+            current_tag_names.append(tag.text)
+            self.set_tags(current_tag_names)
+        
+        return self
+
+    def remove_tag(self, tag: EnhancedResourceTag):
+        """Remove a single enhanced resource tag from the project.
+        
+        Args:
+            tag: EnhancedResourceTag instance to remove
+            
+        Returns:
+            Self for method chaining
+        """
+        current_tags = self.get_tags()
+        current_tag_names = [t.text for t in current_tags if t.uid != tag.uid]
+        self.set_tags(current_tag_names)
+        return self
+
 
 class AlignerrWorkspace:
     def __init__(self, client: "Client"):

libs/labelbox/src/labelbox/alignerr/alignerr_project_builder.py

@@ -4,8 +4,8 @@
 
 from labelbox.alignerr.schema.project_rate import BillingMode
 from labelbox.alignerr.schema.project_rate import ProjectRateInput
-from labelbox.alignerr.schema.project_rate import ProjectRateV2
 from labelbox.alignerr.schema.project_domain import ProjectDomain
+from labelbox.alignerr.schema.enchanced_resource_tags import EnhancedResourceTag, ResourceTagType
 from labelbox.schema.media_type import MediaType
 
 logger = logging.getLogger(__name__)
@@ -22,6 +22,7 @@ def __init__(self, client: "Client"):
         self._alignerr_rates: dict[str, ProjectRateInput] = {}
         self._customer_rate: ProjectRateInput = None
         self._domains: list[ProjectDomain] = []
+        self._enhanced_resource_tags: list[EnhancedResourceTag] = []
         self.role_name_to_id = self._get_role_name_to_id()
 
     def set_name(self, name: str):
@@ -110,6 +111,37 @@ def set_domains(self, domains: list[str]):
             self._domains.append(domain_result)
         return self
 
+    def set_tags(self, tag_texts: list[str], tag_type: ResourceTagType):
+        """Set enhanced resource tags for the project.
+        
+        Args:
+            tag_texts: List of tag text values to search for and attach
+            tag_type: Type filter for searching tags
+            
+        Returns:
+            Self for method chaining
+        """
+        for tag_text in tag_texts:
+            # Search for existing tags by text
+            existing_tags = EnhancedResourceTag.search_by_text(
+                self.client, search_text=tag_text, tag_type=tag_type
+            )
+            
+            if existing_tags:
+                # Use the first matching tag
+                self._enhanced_resource_tags.append(existing_tags[0])
+            else:
+                # Create new tag if not found
+                new_tag = EnhancedResourceTag.create(
+                    self.client,
+                    text=tag_text,
+                    color="#007bff",  # Default blue color
+                    tag_type=tag_type
+                )
+                self._enhanced_resource_tags.append(new_tag)
+        return self
+
+
     def create(self, skip_validation: bool = False):
         if not skip_validation:
             self._validate()
@@ -130,6 +162,7 @@ def create(self, skip_validation: bool = False):
 
         self._create_rates(alignerr_project)
         self._create_domains(alignerr_project)
+        self._create_resource_tags(alignerr_project)
 
         return alignerr_project
 
@@ -150,6 +183,25 @@ def _create_domains(self, alignerr_project: "AlignerrProject"):
                 domain_ids=domain_ids,
             )
 
+    def _create_resource_tags(self, alignerr_project: "AlignerrProject"):
+        if self._enhanced_resource_tags:
+            logger.info(
+                f"Setting enhanced resource tags: {[tag.text for tag in self._enhanced_resource_tags]}"
+            )
+            # Group tags by type and set them accordingly
+            tags_by_type = {}
+            for tag in self._enhanced_resource_tags:
+                tag_type = tag.type
+                if tag_type not in tags_by_type:
+                    tags_by_type[tag_type] = []
+                tags_by_type[tag_type].append(tag.text)
+            
+            # Set tags for each type
+            for tag_type_str, tag_names in tags_by_type.items():
+                # Convert string back to enum
+                tag_type_enum = ResourceTagType(tag_type_str)
+                alignerr_project.set_tags(tag_names, tag_type_enum)
+
     def _validate_alignerr_rates(self):
         # Import here to avoid circular imports
         from labelbox.alignerr.alignerr_project import AlignerrRole

libs/labelbox/src/labelbox/alignerr/schema/enchanced_resource_tags.py

@@ -0,0 +1,139 @@
+from enum import Enum
+from typing import List, Optional
+from labelbox.orm.db_object import DbObject, Updateable
+from labelbox.orm.model import Field
+from pydantic import BaseModel
+
+
+class ResourceTagType(Enum):
+    """Enum for resource tag types."""
+    Default = "Default"
+    System = "System"
+    Request = "Request"
+    Migration = "Migration"
+    Billing = "Billing"
+
+
+class CreateResourceTagInput(BaseModel):
+    """Input for creating a new resource tag."""
+
+    text: str
+    color: str
+    type: Optional[str] = None
+
+
+class UpdateResourceTagInput(BaseModel):
+    """Input for updating a resource tag."""
+
+    id: str
+    text: str
+    color: str
+    type: Optional[str] = None
+
+
+class DeleteResourceTagInput(BaseModel):
+    """Input for deleting a resource tag."""
+
+    id: str
+    type: Optional[str] = None
+
+
+class ResourceTagsInput(BaseModel):
+    """Input for querying resource tags."""
+
+    type: str
+
+
+class EnhancedResourceTag(DbObject, Updateable):
+    """Enhanced resource tag with additional functionality and type support."""
+
+    # Fields matching the DDL schema
+    id = Field.String("id")
+    createdAt = Field.DateTime("createdAt")
+    updatedAt = Field.DateTime("updatedAt")
+    organizationId = Field.String("organizationId")
+    text = Field.String("text")
+    color = Field.String("color")
+    createdById = Field.String("createdById")
+    type = Field.String("type")
+
+    @classmethod
+    def create(
+        cls, client, text: str, color: str, tag_type: Optional[ResourceTagType] = None
+    ) -> "EnhancedResourceTag":
+        """Create a new enhanced resource tag.
+
+        Args:
+            client: Labelbox client instance
+            text: Text content of the resource tag
+            color: Color of the resource tag
+            tag_type: Optional type of the resource tag
+
+        Returns:
+            Created EnhancedResourceTag instance
+        """
+        # Use the existing organization create_resource_tag method
+        # Get the organization
+        org = client.get_organization()
+        
+        # Create the tag using existing API
+        tag_data = {"text": text, "color": color}
+        created_tag = org.create_resource_tag(tag_data)
+        
+        # Create EnhancedResourceTag with the same data plus defaults for missing fields
+        enhanced_tag = cls(client, {
+            "id": created_tag.uid,
+            "text": created_tag.text,
+            "color": created_tag.color,
+            "createdAt": None,
+            "updatedAt": None,
+            "organizationId": None,
+            "createdById": None,
+            "type": tag_type.value if tag_type else None
+        })
+        
+        return enhanced_tag
+
+
+
+    @classmethod
+    def search_by_text(
+        cls, client, search_text: str, tag_type: ResourceTagType
+    ) -> List["EnhancedResourceTag"]:
+        """Search resource tags by text content.
+
+        Args:
+            client: Labelbox client instance
+            search_text: Text to search for
+            tag_type: Type filter
+
+        Returns:
+            List of matching EnhancedResourceTag instances
+        """
+        # Use the existing organization get_resource_tags method
+        # Get the organization
+        org = client.get_organization()
+        
+        # Get all resource tags
+        regular_tags = org.get_resource_tags()
+        
+        # Convert to EnhancedResourceTag instances and filter by search text and type
+        matching_tags = []
+        for tag in regular_tags:
+            if search_text.lower() in tag.text.lower():
+                enhanced_tag = cls(client, {
+                    "id": tag.uid,
+                    "text": tag.text,
+                    "color": tag.color,
+                    "createdAt": None,
+                    "updatedAt": None,
+                    "organizationId": None,
+                    "createdById": None,
+                    "type": tag_type.value
+                })
+                
+                # Apply type filter
+                if enhanced_tag.type == tag_type.value:
+                    matching_tags.append(enhanced_tag)
+        
+        return matching_tags

libs/labelbox/src/labelbox/alignerr/schema/project_rate.py

@@ -53,7 +53,7 @@ class ProjectRateV2(DbObject, Deletable):
     effectiveUntil = Field.DateTime("effectiveUntil")
 
     @classmethod
-    def get_by_project_id(cls, client, project_id: str) -> "ProjectRateV2":
+    def get_by_project_id(cls, client, project_id: str) -> list["ProjectRateV2"]:
         query_str = """
         query GetAllProjectRatesPyApi($projectId: ID!) {
             project(where: { id: $projectId }) {
@@ -84,10 +84,10 @@ def get_by_project_id(cls, client, project_id: str) -> "ProjectRateV2":
         rates_data = result["project"]["ratesV2"]
 
         if not rates_data:
-            return None
+            return []
 
-        # Return the first rate as a ProjectRateV2 object
-        return cls(client, rates_data[0])
+        # Return all rates as ProjectRateV2 objects
+        return [cls(client, rate_data) for rate_data in rates_data]
 
     @classmethod
     def set_project_rate(

libs/labelbox/tests/integration/test_alignerr_project.py

@@ -73,16 +73,16 @@ def test_alignerr_project_domains(client, test_alignerr_project):
     # The collection might be empty for a new project, which is expected
 
 
-def test_alignerr_project_get_project_rate_no_rates(
+def test_alignerr_project_get_project_rates_no_rates(
     client, test_alignerr_project
 ):
-    """Test get_project_rate() when no rates are set."""
-    # For a new project without rates, this should return None
-    project_rate = test_alignerr_project.get_project_rate()
-    assert project_rate is None
+    """Test get_project_rates() when no rates are set."""
+    # For a new project without rates, this should return an empty list
+    project_rates = test_alignerr_project.get_project_rates()
+    assert project_rates == []
 
 
-def test_alignerr_project_set_and_get_project_rate(
+def test_alignerr_project_set_and_get_project_rates(
     client, test_alignerr_project
 ):
     """Test setting and getting project rates."""
@@ -100,7 +100,10 @@ def test_alignerr_project_set_and_get_project_rate(
     result = test_alignerr_project.set_project_rate(project_rate_input)
     assert result is True  # Should return success status
 
-    # Get the project rate back
-    project_rate = test_alignerr_project.get_project_rate()
+    # Get the project rates back
+    project_rates = test_alignerr_project.get_project_rates()
+    # Should return a list with at least one rate
+    assert isinstance(project_rates, list)
+    assert len(project_rates) >= 1
     # Note: The actual rate retrieval might depend on the API implementation
     # This test verifies the method calls work without errors