Made the required changes

Author: vaibhavatlan

pyatlan/model/assets/core/a_i_model.py

@@ -4,21 +4,16 @@
 
 from __future__ import annotations
 
-import sys
 from typing import ClassVar, Dict, List, Optional, Set, overload
 
 from pydantic.v1 import Field, validator
 
 from pyatlan.model.enums import AIDatasetType, AIModelStatus, AtlanConnectorType
 from pyatlan.model.fields.atlan_fields import KeywordField, RelationField, TextField
-from pyatlan.utils import (
-    get_epoch_timestamp,
-    init_guid,
-    to_camel_case,
-    validate_required_fields,
-)
+from pyatlan.utils import init_guid, to_camel_case, validate_required_fields
 
 from .a_i import AI
+from .asset import Asset
 from .process import Process
 
 
@@ -73,34 +68,33 @@ def creator(
     @classmethod
     def processes_creator(
         cls,
-        a_i_model_guid: str,
-        a_i_model_name: str,
-        database_dict: Dict[AIDatasetType, list],
+        ai_model: AIModel,
+        dataset_dict: Dict[AIDatasetType, list],
     ) -> List[Process]:
+        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 database_dict.items():
+        for key, value_list in dataset_dict.items():
             for value in value_list:
-                asset_type = getattr(
-                    sys.modules.get("pyatlan.model.assets", {}), value.type_name, None
-                )
+                asset_type = Asset._convert_to_real_type_(value)
                 if key == AIDatasetType.OUTPUT:
-                    process_name = f"{a_i_model_name} -> {value.name}"
+                    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=a_i_model_guid)],
+                        inputs=[AIModel.ref_by_guid(guid=ai_model.guid)],
                         outputs=[asset_type.ref_by_guid(guid=value.guid)],  # type: ignore
-                        process_id=str(get_epoch_timestamp()),
+                        extra_hash_params={key.value},
                     )
                     process_created.ai_dataset_type = key
                 else:
-                    process_name = f"{value.name} -> {a_i_model_name}"
+                    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=a_i_model_guid)],
-                        process_id=str(get_epoch_timestamp()),
+                        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)

pyatlan/model/assets/core/process.py

@@ -31,6 +31,7 @@ def creator(
         outputs: List["Catalog"],
         process_id: Optional[str] = None,
         parent: Optional[Process] = None,
+        extra_hash_params: Optional[Set[str]] = set(),
     ) -> Process:
         return Process(
             attributes=Process.Attributes.create(
@@ -40,6 +41,7 @@ def creator(
                 inputs=inputs,
                 outputs=outputs,
                 parent=parent,
+                extra_hash_params=extra_hash_params,
             )
         )
 
@@ -383,6 +385,7 @@ def generate_qualified_name(
             outputs: List["Catalog"],
             parent: Optional["Process"] = None,
             process_id: Optional[str] = None,
+            extra_hash_params: Optional[Set[str]] = set(),
         ) -> str:
             def append_relationship(output: StringIO, relationship: Asset):
                 if relationship.guid:
@@ -405,6 +408,11 @@ def append_relationships(output: StringIO, relationships: List["Catalog"]):
                 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()
@@ -421,6 +429,7 @@ def create(
             outputs: List["Catalog"],
             process_id: Optional[str] = None,
             parent: Optional[Process] = None,
+            extra_hash_params: Optional[Set[str]] = set(),
         ) -> Process.Attributes:
             qualified_name = Process.Attributes.generate_qualified_name(
                 name=name,
@@ -429,6 +438,7 @@ def create(
                 inputs=inputs,
                 outputs=outputs,
                 parent=parent,
+                extra_hash_params=extra_hash_params,
             )
             connector_name = connection_qualified_name.split("/")[1]
             return Process.Attributes(

tests/integration/ai_asset_test.py

@@ -115,6 +115,47 @@ def test_update_ai_assets(
     _update_ai_model(client, ai_model)
 
 
+def _assert_response_processes_creator(
+    mutation_response, asset_list, ai_dataset_type, process_sum, ai_model
+):
+    for i in range(len(asset_list)):
+        assert mutation_response.mutated_entities.CREATE[i + process_sum]
+        assert (
+            mutation_response.mutated_entities.CREATE[i + process_sum].ai_dataset_type  # type: ignore
+            == ai_dataset_type
+        )
+        if ai_dataset_type == AIDatasetType.OUTPUT:
+            assert (
+                mutation_response.mutated_entities.CREATE[i + process_sum].inputs  # type: ignore
+                and mutation_response.mutated_entities.CREATE[i + process_sum]
+                .inputs[0]
+                .guid
+                == ai_model.guid  # type: ignore
+            )
+            assert (
+                mutation_response.mutated_entities.CREATE[i + process_sum].outputs  # type: ignore
+                and mutation_response.mutated_entities.CREATE[i + process_sum]
+                .outputs[0]
+                .guid  # type: ignore
+                == asset_list[i].guid
+            )
+        else:
+            assert (
+                mutation_response.mutated_entities.CREATE[i + process_sum].inputs  # type: ignore
+                and mutation_response.mutated_entities.CREATE[i + process_sum]
+                .inputs[0]
+                .guid
+                == asset_list[i].guid  # type: ignore
+            )
+            assert (
+                mutation_response.mutated_entities.CREATE[i + process_sum].outputs  # type: ignore
+                and mutation_response.mutated_entities.CREATE[i + process_sum]
+                .outputs[0]
+                .guid  # type: ignore
+                == ai_model.guid
+            )
+
+
 def test_ai_model_processes_creator(
     client: AtlanClient,
     ai_model: AIModel,
@@ -159,17 +200,16 @@ def test_ai_model_processes_creator(
         list_validation.append(results)
         list_output.append(results)
 
-    database_dict = {
+    dataset_dict = {
         AIDatasetType.TRAINING: list_training,
         AIDatasetType.TESTING: list_testing,
         AIDatasetType.INFERENCE: list_inference,
         AIDatasetType.VALIDATION: list_validation,
         AIDatasetType.OUTPUT: list_output,
     }
     created_processes = AIModel.processes_creator(
-        a_i_model_guid=ai_model.guid,
-        a_i_model_name=AI_MODEL_NAME,  # Add fallback for type safety
-        database_dict=database_dict,
+        ai_model=ai_model,
+        dataset_dict=dataset_dict,
     )
     response = AIModel.processes_batch_save(client, created_processes)
 
@@ -178,111 +218,42 @@ def test_ai_model_processes_creator(
     assert (
         mutation_response.mutated_entities and mutation_response.mutated_entities.CREATE
     )
-    for i in range(len(list_training)):
-        assert mutation_response.mutated_entities.CREATE[i]
-        assert (
-            mutation_response.mutated_entities.CREATE[i].ai_dataset_type  # type: ignore
-            == AIDatasetType.TRAINING
-        )
-        assert (
-            mutation_response.mutated_entities.CREATE[i].inputs  # type: ignore
-            and mutation_response.mutated_entities.CREATE[i].inputs[0].guid
-            == list_training[i].guid  # type: ignore
-        )
-        assert (
-            mutation_response.mutated_entities.CREATE[i].outputs  # type: ignore
-            and mutation_response.mutated_entities.CREATE[i].outputs[0].guid  # type: ignore
-            == ai_model.guid
-        )
-    current_process_sum = len(list_training)
-    for i in range(len(list_testing)):
-        assert mutation_response.mutated_entities.CREATE[i + current_process_sum]
-        assert (
-            mutation_response.mutated_entities.CREATE[
-                i + current_process_sum
-            ].ai_dataset_type  # type: ignore
-            == AIDatasetType.TESTING
-        )
-        assert (
-            mutation_response.mutated_entities.CREATE[i + current_process_sum].inputs  # type: ignore
-            and mutation_response.mutated_entities.CREATE[i + current_process_sum]
-            .inputs[0]
-            .guid
-            == list_testing[i].guid  # type: ignore
-        )
-        assert (
-            mutation_response.mutated_entities.CREATE[i + current_process_sum].outputs  # type: ignore
-            and mutation_response.mutated_entities.CREATE[i + current_process_sum]
-            .outputs[0]
-            .guid  # type: ignore
-            == ai_model.guid
-        )
-    current_process_sum += len(list_testing)
-    for i in range(len(list_inference)):
-        assert mutation_response.mutated_entities.CREATE[i + current_process_sum]
-        assert (
-            mutation_response.mutated_entities.CREATE[
-                i + current_process_sum
-            ].ai_dataset_type  # type: ignore
-            == AIDatasetType.INFERENCE
-        )
-        assert (
-            mutation_response.mutated_entities.CREATE[i + current_process_sum].inputs  # type: ignore
-            and mutation_response.mutated_entities.CREATE[i + current_process_sum]
-            .inputs[0]
-            .guid
-            == list_inference[i].guid  # type: ignore
-        )
-        assert (
-            mutation_response.mutated_entities.CREATE[i + current_process_sum].outputs  # type: ignore
-            and mutation_response.mutated_entities.CREATE[i + current_process_sum]
-            .outputs[0]
-            .guid  # type: ignore
-            == ai_model.guid
-        )
-    current_process_sum += len(list_inference)
-    for i in range(len(list_validation)):
-        assert mutation_response.mutated_entities.CREATE[i + current_process_sum]
-        assert (
-            mutation_response.mutated_entities.CREATE[
-                i + current_process_sum
-            ].ai_dataset_type  # type: ignore
-            == AIDatasetType.VALIDATION
-        )
-        assert (
-            mutation_response.mutated_entities.CREATE[i + current_process_sum].inputs  # type: ignore
-            and mutation_response.mutated_entities.CREATE[i + current_process_sum]
-            .inputs[0]
-            .guid
-            == list_validation[i].guid  # type: ignore
-        )
-        assert (
-            mutation_response.mutated_entities.CREATE[i + current_process_sum].outputs  # type: ignore
-            and mutation_response.mutated_entities.CREATE[i + current_process_sum]
-            .outputs[0]
-            .guid  # type: ignore
-            == ai_model.guid
-        )
-    current_process_sum += len(list_validation)
-    for i in range(len(list_output)):
-        assert mutation_response.mutated_entities.CREATE[i + current_process_sum]
-        assert (
-            mutation_response.mutated_entities.CREATE[
-                i + current_process_sum
-            ].ai_dataset_type  # type: ignore
-            == AIDatasetType.OUTPUT
-        )
-        assert (
-            mutation_response.mutated_entities.CREATE[i + current_process_sum].inputs  # type: ignore
-            and mutation_response.mutated_entities.CREATE[i + current_process_sum]
-            .inputs[0]
-            .guid
-            == ai_model.guid  # type: ignore
-        )
-        assert (
-            mutation_response.mutated_entities.CREATE[i + current_process_sum].outputs  # type: ignore
-            and mutation_response.mutated_entities.CREATE[i + current_process_sum]
-            .outputs[0]
-            .guid  # type: ignore
-            == list_output[i].guid
-        )
+    currnt_processes_sum = 0
+    _assert_response_processes_creator(
+        mutation_response, list_training, AIDatasetType.TRAINING, 0, ai_model
+    )
+    currnt_processes_sum += len(list_training)
+    _assert_response_processes_creator(
+        mutation_response,
+        list_testing,
+        AIDatasetType.TESTING,
+        currnt_processes_sum,
+        ai_model,
+    )
+    currnt_processes_sum += len(list_testing)
+    _assert_response_processes_creator(
+        mutation_response,
+        list_inference,
+        AIDatasetType.INFERENCE,
+        currnt_processes_sum,
+        ai_model,
+    )
+    currnt_processes_sum += len(list_inference)
+    _assert_response_processes_creator(
+        mutation_response,
+        list_validation,
+        AIDatasetType.VALIDATION,
+        currnt_processes_sum,
+        ai_model,
+    )
+    currnt_processes_sum += len(list_validation)
+    _assert_response_processes_creator(
+        mutation_response,
+        list_output,
+        AIDatasetType.OUTPUT,
+        currnt_processes_sum,
+        ai_model,
+    )
+    currnt_processes_sum += len(list_output)
+
+    assert currnt_processes_sum == len(created_processes)