Merge branch 'main' into APP-7824

Author: Aryamanz29

pyatlan/generator/templates/imports.jinja2

@@ -120,6 +120,7 @@ from pyatlan.model.structs import (
 from pyatlan.utils import (
     init_guid,
     next_id,
+    to_camel_case,
     validate_required_fields,
     validate_single_required_field,
     get_parent_qualified_name,

pyatlan/generator/templates/methods/asset/a_i_application.jinja2

@@ -0,0 +1,49 @@
+
+    @overload
+    @classmethod
+    def creator(
+        cls,
+        *,
+        name: str,
+        ai_application_version: str,
+        ai_application_development_stage: AIApplicationDevelopmentStage,
+    ) -> AIApplication: ...
+
+    @overload
+    @classmethod
+    def creator(
+        cls,
+        *,
+        name: str,
+        ai_application_version: str,
+        ai_application_development_stage: AIApplicationDevelopmentStage,
+        owner_groups: Set[str],
+        owner_users: Set[str],
+        models: List[AIModel],
+    ) -> AIApplication: ...
+
+    @classmethod
+    @init_guid
+    def creator(
+        cls,
+        *,
+        name: str,
+        ai_application_version: str,
+        ai_application_development_stage: AIApplicationDevelopmentStage,
+        owner_groups: Optional[Set[str]] = None,
+        owner_users: Optional[Set[str]] = None,
+        models: Optional[List[AIModel]] = None,
+    ) -> AIApplication:
+        validate_required_fields(
+            ["name", "ai_application_version", "ai_application_development_stage"],
+            [name, ai_application_version, ai_application_development_stage],
+        )
+        attributes = AIApplication.Attributes.creator(
+            name=name,
+            ai_application_version=ai_application_version,
+            ai_application_development_stage=ai_application_development_stage,
+            owner_groups=owner_groups,
+            owner_users=owner_users,
+            models=models,
+        )
+        return cls(attributes=attributes)

pyatlan/generator/templates/methods/asset/a_i_model.jinja2

@@ -0,0 +1,110 @@
+
+    @overload
+    @classmethod
+    def creator(
+        cls,
+        *,
+        name: str,
+        ai_model_status: AIModelStatus,
+    ) -> AIModel: ...
+
+    @overload
+    @classmethod
+    def creator(
+        cls,
+        *,
+        name: str,
+        ai_model_status: AIModelStatus,
+        owner_groups: Set[str],
+        owner_users: Set[str],
+        ai_model_version: str,
+    ) -> AIModel: ...
+
+    @classmethod
+    @init_guid
+    def creator(
+        cls,
+        *,
+        name: str,
+        ai_model_status: AIModelStatus,
+        owner_groups: Optional[Set[str]] = None,
+        owner_users: Optional[Set[str]] = None,
+        ai_model_version: Optional[str] = None,
+    ) -> AIModel:
+        validate_required_fields(
+            ["name", "ai_model_status"],
+            [name, ai_model_status],
+        )
+        attributes = AIModel.Attributes.creator(
+            name=name,
+            ai_model_status=ai_model_status,
+            owner_groups=owner_groups,
+            owner_users=owner_users,
+            ai_model_version=ai_model_version,
+        )
+        return cls(attributes=attributes)
+
+    @classmethod
+    def processes_creator(
+        cls,
+        ai_model: AIModel,
+        dataset_dict: Dict[AIDatasetType, list],
+    ) -> List[Process]:
+        """
+        Creates a list of Process objects representing the relationships between an AI model and its datasets.
+
+        :param ai_model: the AI model for which to create processes
+        :param dataset_dict: dictionary mapping AI dataset types to lists of assets
+        :returns: list of Process objects representing the AI model's data lineage
+        :raises ValueError: when the AI model is missing required attributes (guid or name)
+        """
+        if not ai_model.guid or not ai_model.name:
+            raise ValueError("AI model must have both guid and name attributes")
+        process_list = []
+        for key, value_list in dataset_dict.items():
+            for value in value_list:
+                asset_type = Asset._convert_to_real_type_(value)
+                if key == AIDatasetType.OUTPUT:
+                    process_name = f"{ai_model.name} -> {value.name}"
+                    process_created = Process.creator(
+                        name=process_name,
+                        connection_qualified_name="default/ai/dataset",
+                        inputs=[AIModel.ref_by_guid(guid=ai_model.guid)],
+                        outputs=[asset_type.ref_by_guid(guid=value.guid)],  # type: ignore
+                        extra_hash_params={key.value},
+                    )
+                    process_created.ai_dataset_type = key
+                else:
+                    process_name = f"{value.name} -> {ai_model.name}"
+                    process_created = Process.creator(
+                        name=process_name,
+                        connection_qualified_name="default/ai/dataset",
+                        inputs=[asset_type.ref_by_guid(guid=value.guid)],  # type: ignore
+                        outputs=[AIModel.ref_by_guid(guid=ai_model.guid)],
+                        extra_hash_params={key.value},
+                    )
+                    process_created.ai_dataset_type = key
+                process_list.append(process_created)
+
+        return process_list
+
+    @classmethod
+    def processes_batch_save(cls, client, process_list: List[Process]) -> List:
+        """
+        Saves a list of Process objects to Atlan in batches to optimize performance.
+        We save the processes in batches of 20.
+
+        :param client: Atlan client instance for making API calls
+        :param process_list: list of Process objects to save
+        :returns: list of API responses from each batch save operation
+        """
+        batch_size = 20
+        total_processes = len(process_list)
+        responses = []
+
+        for i in range(0, total_processes, batch_size):
+            batch = process_list[i : i + batch_size]
+            response = client.asset.save(batch)
+            responses.append(response)
+
+        return responses

pyatlan/generator/templates/methods/asset/process.jinja2

@@ -9,6 +9,7 @@
         outputs: List["Catalog"],
         process_id: Optional[str] = None,
         parent: Optional[Process] = None,
+        extra_hash_params: Optional[Set[str]] = None,
     ) -> Process:
         return Process(
             attributes=Process.Attributes.create(
@@ -18,6 +19,7 @@
                 inputs=inputs,
                 outputs=outputs,
                 parent=parent,
+                extra_hash_params=extra_hash_params,
             )
         )
 
@@ -31,6 +33,7 @@
         outputs: List["Catalog"],
         process_id: Optional[str] = None,
         parent: Optional[Process] = None,
+        extra_hash_params: Optional[Set[str]] = None,
     ) -> Process:
         warn(
             (
@@ -47,4 +50,5 @@
             outputs=outputs,
             process_id=process_id,
             parent=parent,
+            extra_hash_params=extra_hash_params,
         )

pyatlan/generator/templates/methods/attribute/a_i_application.jinja2

@@ -0,0 +1,33 @@
+
+        @classmethod
+        @init_guid
+        def creator(
+            cls,
+            *,
+            name: str,
+            ai_application_version: str,
+            ai_application_development_stage: AIApplicationDevelopmentStage,
+            owner_groups: Optional[Set[str]] = None,
+            owner_users: Optional[Set[str]] = None,
+            models: Optional[List[AIModel]] = None,
+        ) -> AIApplication.Attributes:
+            validate_required_fields(
+                ["name", "ai_application_version", "ai_application_development_stage"],
+                [name, ai_application_version, ai_application_development_stage],
+            )
+            owner_groups = owner_groups or set()
+            owner_users = owner_users or set()
+            models = models or []
+            name_camel_case = to_camel_case(name)
+            return AIApplication.Attributes(
+                name=name,
+                qualified_name=f"default/ai/aiapplication/{name_camel_case}",
+                connector_name=AtlanConnectorType.AI.value,
+                ai_application_version=ai_application_version,
+                ai_application_development_stage=ai_application_development_stage,
+                owner_groups=owner_groups,
+                owner_users=owner_users,
+                models=models,
+                certificate_status=CertificateStatus.DRAFT,
+                asset_cover_image="/assets/default-product-cover-DeQonY47.webp",
+            )

pyatlan/generator/templates/methods/attribute/a_i_model.jinja2

@@ -0,0 +1,29 @@
+
+        @classmethod
+        @init_guid
+        def creator(
+            cls,
+            *,
+            name: str,
+            ai_model_status: AIModelStatus,
+            owner_groups: Optional[Set[str]] = None,
+            owner_users: Optional[Set[str]] = None,
+            ai_model_version: Optional[str] = None,
+        ) -> AIModel.Attributes:
+            validate_required_fields(
+                ["name", "ai_model_status"],
+                [name, ai_model_status],
+            )
+            owner_groups = owner_groups or set()
+            owner_users = owner_users or set()
+            name_camel_case = to_camel_case(name)
+            return AIModel.Attributes(
+                name=name,
+                qualified_name=f"default/ai/aiapplication/{name_camel_case}",
+                connector_name=AtlanConnectorType.AI.value,
+                ai_model_status=ai_model_status,
+                ai_model_version=ai_model_version,
+                owner_groups=owner_groups,
+                owner_users=owner_users,
+                asset_cover_image="/assets/default-product-cover-DeQonY47.webp",
+            )

pyatlan/generator/templates/methods/attribute/process.jinja2

@@ -7,6 +7,7 @@
             outputs: List["Catalog"],
             parent: Optional["Process"] = None,
             process_id: Optional[str] = None,
+            extra_hash_params: Optional[Set[str]] = None,
         ) -> str:
             def append_relationship(output: StringIO, relationship: Asset):
                 if relationship.guid:
@@ -20,6 +21,7 @@
                 ["name", "connection_qualified_name", "inputs", "outputs"],
                 [name, connection_qualified_name, inputs, outputs],
             )
+            extra_hash_params = extra_hash_params or set()
             if process_id and process_id.strip():
                 return f"{connection_qualified_name}/{process_id}"
             buffer = StringIO()
@@ -29,6 +31,11 @@
                 append_relationship(buffer, parent)
             append_relationships(buffer, inputs)
             append_relationships(buffer, outputs)
+            # Handles edge case where identical name, connection, input, and output caused hash collisions,
+            # resulting in duplicate qualified names and backend skipping process creation.
+            if extra_hash_params:
+                for param in extra_hash_params:
+                    buffer.write(param)
             ret_value = hashlib.md5(  # noqa: S303, S324
                 buffer.getvalue().encode()
             ).hexdigest()
@@ -45,6 +52,7 @@
             outputs: List["Catalog"],
             process_id: Optional[str] = None,
             parent: Optional[Process] = None,
+            extra_hash_params: Optional[Set[str]] = None,
         ) -> Process.Attributes:
             qualified_name = Process.Attributes.generate_qualified_name(
                 name=name,
@@ -53,6 +61,7 @@
                 inputs=inputs,
                 outputs=outputs,
                 parent=parent,
+                extra_hash_params=extra_hash_params,
             )
             connector_name = connection_qualified_name.split("/")[1]
             return Process.Attributes(

pyatlan/generator/templates/module.jinja2

@@ -20,6 +20,11 @@ from pyatlan.model.enums import (
 from .asset import SelfAsset
 {% endif %}
 
+{% if asset_info.name == 'AIModel' %}
+from .asset import Asset
+from .process import Process
+{% endif %}
+
 {% set entity_def = asset_info.entity_def %}
 
 {% set file_name = 'methods/imports/' + entity_def.name | to_snake_case + '.jinja2' %}

pyatlan/model/assets/__init__.py

@@ -88,8 +88,8 @@
         "SnowflakeStage",
         "DatabricksUnityCatalogTag",
         "SnowflakeStream",
-        "Database",
         "CalculationView",
+        "Database",
         "Procedure",
         "SnowflakeTag",
         "MatillionGroup",

pyatlan/model/assets/__init__.pyi

@@ -85,8 +85,8 @@ __all__ = [
     "SnowflakeStage",
     "DatabricksUnityCatalogTag",
     "SnowflakeStream",
-    "Database",
     "CalculationView",
+    "Database",
     "Procedure",
     "SnowflakeTag",
     "MatillionGroup",