APP-7827: Added support for AttributeDef.Options.applicable_ai_asset_types

Author: Aryamanz29

pyatlan/model/constants.py

@@ -14,6 +14,7 @@
         "AtlasGlossaryTerm",
     ]
 ]
+AIAssetTypes = Set[Literal["AIApplication", "AIModel"]]
 OtherAssetTypes = Set[Literal["File"]]
 AssetTypes = Set[
     Literal[

pyatlan/model/typedef.py

@@ -21,6 +21,7 @@
 from pyatlan.errors import ErrorCode
 from pyatlan.model.atlan_image import AtlanImage
 from pyatlan.model.constants import (
+    AIAssetTypes,
     AssetTypes,
     DomainTypes,
     EntityTypes,
@@ -192,6 +193,7 @@
     "DataProduct",
 }
 
+_all_ai_asset_types: AIAssetTypes = {"AIApplication", "AIModel"}
 _all_other_types: OtherAssetTypes = {"File"}
 
 
@@ -494,6 +496,12 @@ class Options(AtlanObject):
             "These cover asset types in data products and data domains. "
             "Only assets of one of these types will have this attribute available.",
         )
+        applicable_ai_asset_types: Optional[str] = Field(
+            default=None,
+            alias="aiAssetsTypeList",
+            description="AI asset type names to which this attribute is restricted. "
+            "Only AI assets of the specified types will have this attribute available.",
+        )
         applicable_other_asset_types: Optional[str] = Field(
             default=None,
             alias="otherAssetTypeList",
@@ -639,6 +647,7 @@ def __setattr__(self, name, value):
         "applicable_connections",
         "applicable_glossaries",
         "applicable_domains",
+        "applicable_ai_asset_types",
     ]
 
     @property
@@ -747,6 +756,30 @@ def applicable_domain_types(self, domain_types: DomainTypes):
             )
         self.options.applicable_domain_types = json.dumps(list(domain_types))
 
+    @property
+    def applicable_ai_asset_types(self) -> AIAssetTypes:
+        """
+        AI asset type names to which this attribute is restricted.
+        Only AI assets of the specified types will have this attribute available.
+        """
+        if self.options and self.options.applicable_ai_asset_types:
+            return set(json.loads(self.options.applicable_ai_asset_types))
+        return set()
+
+    @applicable_ai_asset_types.setter
+    def applicable_ai_asset_types(self, ai_asset_types: AIAssetTypes):
+        if self.options is None:
+            raise ErrorCode.MISSING_OPTIONS.exception_with_parameters()
+        if not isinstance(ai_asset_types, set):
+            raise ErrorCode.INVALID_PARAMETER_TYPE.exception_with_parameters(
+                "applicable_ai_asset_types", AIAssetTypes
+            )
+        if not ai_asset_types.issubset(_all_ai_asset_types):
+            raise ErrorCode.INVALID_PARAMETER_VALUE.exception_with_parameters(
+                ai_asset_types, "applicable_ai_asset_types", _all_domain_types
+            )
+        self.options.applicable_ai_asset_types = json.dumps(list(ai_asset_types))
+
     @property
     def applicable_other_asset_types(self) -> OtherAssetTypes:
         """
@@ -855,6 +888,7 @@ def create(
         applicable_other_asset_types: Optional[OtherAssetTypes] = None,
         applicable_domains: Optional[Set[str]] = None,
         applicable_domain_types: Optional[DomainTypes] = None,
+        applicable_ai_asset_types: Optional[AIAssetTypes] = None,
         description: Optional[str] = None,
     ) -> AttributeDef:
         from pyatlan.utils import validate_required_fields
@@ -918,6 +952,7 @@ def create(
             applicable_glossaries or _get_all_qualified_names(client, "AtlasGlossary")
         )
         attr_def.applicable_domains = applicable_domains or _all_domains
+        attr_def.applicable_ai_asset_types = applicable_ai_asset_types or set()
         return attr_def
 
     def is_archived(self) -> bool:

tests/unit/test_typedef_model.py

@@ -424,6 +424,7 @@ def test_attribute_create_with_limited_applicability(self, client: AtlanClient):
                 "default/snowflake/16992681799",
             },
             applicable_domains={"default/domain/uuBI8WSqeom1PXs7oo20L/super"},
+            applicable_ai_asset_types={"AIModel", "AIApplication"},
         )
         attribute_def_with_limited = AttributeDef.create(
             client=client,