Merge pull request #686 from atlanhq/APP-7827

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_ai_asset_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/constants.py

@@ -1,5 +1,6 @@
 from pyatlan.model.assets import AtlasGlossary
 from pyatlan.model.constants import (
+    AIAssetTypes,
     AssetTypes,
     DomainTypes,
     EntityTypes,
@@ -599,6 +600,8 @@
 
 APPLICABLE_DOMAIN_TYPES = "applicable_domain_types"
 
+APPLICABLE_AI_ASSET_TYPES = "applicable_ai_asset_types"
+
 APPLICABLE_ASSET_TYPES = "applicable_asset_types"
 
 TEST_ENUM_DEF = {
@@ -721,6 +724,16 @@
         {"Bogus"},
         "ATLAN-PYTHON-400-051 Invalid value(s) {'Bogus'} provided for 'applicable_domain_types'.",
     ),
+    (
+        APPLICABLE_AI_ASSET_TYPES,
+        1,
+        f"ATLAN-PYTHON-400-048 Invalid parameter type for applicable_ai_asset_types should be {AIAssetTypes}",
+    ),
+    (
+        APPLICABLE_AI_ASSET_TYPES,
+        {"Bogus"},
+        "ATLAN-PYTHON-400-051 Invalid value(s) {'Bogus'} provided for 'applicable_ai_asset_types'.",
+    ),
     (
         APPLICABLE_OTHER_ASSET_TYPES,
         1,

tests/unit/test_typedef_model.py

@@ -24,13 +24,15 @@
     StructDef,
     TypeDef,
     TypeDefResponse,
+    _all_ai_asset_types,
     _all_domain_types,
     _all_glossary_types,
     _all_other_types,
     _complete_type_list,
 )
 from pyatlan.model.utils import to_camel_case, to_snake_case
 from tests.unit.constants import (
+    APPLICABLE_AI_ASSET_TYPES,
     APPLICABLE_ASSET_TYPES,
     APPLICABLE_CONNECTIONS,
     APPLICABLE_DOMAIN_TYPES,
@@ -363,6 +365,7 @@ def sut(self, client: AtlanClient) -> AttributeDef:
             (APPLICABLE_ASSET_TYPES, {"Table"}),
             (APPLICABLE_GLOSSARY_TYPES, {"AtlasGlossary"}),
             (APPLICABLE_DOMAIN_TYPES, {"DataDomain", "DataProduct"}),
+            (APPLICABLE_AI_ASSET_TYPES, {"AIApplication", "AIModel"}),
             (APPLICABLE_OTHER_ASSET_TYPES, {"File"}),
             (APPLICABLE_ENTITY_TYPES, {"Asset"}),
         ],
@@ -397,6 +400,7 @@ def test_applicable_types_with_invalid_type_raises_invalid_request_error(
             (APPLICABLE_ASSET_TYPES, {random.choice(list(_complete_type_list))}),
             (APPLICABLE_GLOSSARY_TYPES, {random.choice(list(_all_glossary_types))}),
             (APPLICABLE_DOMAIN_TYPES, {random.choice(list(_all_domain_types))}),
+            (APPLICABLE_AI_ASSET_TYPES, {random.choice(list(_all_ai_asset_types))}),
             (APPLICABLE_OTHER_ASSET_TYPES, {random.choice(list(_all_other_types))}),
             (APPLICABLE_ENTITY_TYPES, {"Asset"}),
             (APPLICABLE_CONNECTIONS, {"default/snowflake/1699268171"}),
@@ -424,6 +428,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,